Repository: CitrineInformatics/lolo (language: python)

Issue:
Consolidate GuessTheMean implementations
I broke these apart previously, but then realized everything besides the `Learner` interface actually is generic and can be implemented with the same generic class. So consolidating them here.

This is the last PR going into this feature branch before a merge to main. I do not immediately plan to cut a release though, because I may look into implementing a multi-task linear regression learner.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index ea8cad88..bc3421e0 100644
--- a/README.md
+++ b/README.md
@@ -43,18 +43,18 @@ For example, you can use Random Forest with:
 import io.citrine.lolo.learners.RandomForest
 val trainingData: Seq[(Vector[Any], Any)] = features.zip(labels)
 val model = new RandomForest().train(trainingData).getModel()
-val predictions: Seq[Any] = model.transform(testInputs).getExpected()
+val predictions: Seq[Any] = model.transform(testInputs).expected
 ```
 
 # Performance
 Lolo prioritizes functionality over performance, but it is still quite fast.  In its _random forest_ use case, the complexity scales as:
 
 | Time complexity | Training rows | Features | Trees |
-|-------|--------|-------|-------|
-| `train` | O(n log n) | O(n) | O(n) |
-| `getLoss` | O(n log n) | O(n) | O(n) |
-| `getExpected` | O(log n) | O(1) | O(n) |
-| `getUncertainty` | O(n) | O(1) | O(n) |
+|-----------------|--------|-------|-------|
+| `train`         | O(n log n) | O(n) | O(n) |
+| `loss`          | O(n log n) | O(n) | O(n) |
+| `expected`      | O(log n) | O(1) | O(n) |
+| `uncertainty`   | O(n) | O(1) | O(n) |
 
 On an [Ivy Bridge](http://ark.intel.com/products/77780/Intel-Core-i7-4930K-Processor-12M-Cache-up-to-3_90-GHz) test platform, the (1024 row, 1024 tree, 8 feature) [performance test](src/test/scala/io/citrine/lolo/PerformanceTest.scala) took 1.4 sec to train and 2.3 ms per prediction with uncertainty.
 
diff --git a/python/examples/profile/scala-benchmark.scala b/python/examples/profile/scala-benchmark.scala
index 850a9d3a..876e78e9 100644
--- a/python/examples/profile/scala-benchmark.scala
+++ b/python/examples/profile/scala-benchmark.scala
@@ -18,12 +18,12 @@ def timedTest(trainingData: Seq[(Vector[Any], Double)], evalData: Seq[(Vector[An
   // Do dry timining run with 8 runs
   val model = baggedLearner.train(trainingData).getModel()
   Stopwatch.time({baggedLearner.train(trainingData).getModel()}, benchmark = "None", minRun = 8, maxRun = 8)
-  Stopwatch.time({model.transform(inputs).getExpected()}, benchmark = "None", minRun = 8, maxRun = 8)
+  Stopwatch.time({model.transform(inputs).expected}, benchmark = "None", minRun = 8, maxRun = 8)
   Stopwatch.time({model.transform(inputs).getUncertainty()}, benchmark = "None", minRun = 8, maxRun = 8)
 
   // Time it for real
   val timeTraining = Stopwatch.time({baggedLearner.train(trainingData).getModel()}, benchmark = "None", minRun = 16, targetError = 0.1, maxRun = 32)
-  val timePredicting = Stopwatch.time({model.transform(inputs).getExpected()}, benchmark = "None", minRun = 16, targetError = 0.1, maxRun = 32)
+  val timePredicting = Stopwatch.time({model.transform(inputs).expected}, benchmark = "None", minRun = 16, targetError = 0.1, maxRun = 32)
   val timeUncertainty = Stopwatch.time({model.transform(inputs).getUncertainty()}, benchmark = "None", minRun = 16, targetError = 0.1, maxRun = 32)
 
   (timeTraining, timePredicting, timeUncertainty)
diff --git a/python/lolopy/learners.py b/python/lolopy/learners.py
index 440670c3..29d906a0 100644
--- a/python/lolopy/learners.py
+++ b/python/lolopy/learners.py
@@ -6,7 +6,13 @@
 from sklearn.base import BaseEstimator, RegressorMixin, ClassifierMixin, is_regressor
 from sklearn.exceptions import NotFittedError
 
-__all__ = ['RandomForestRegressor', 'RandomForestClassifier', 'ExtraRandomTreesRegressor', 'ExtraRandomTreesClassifier']
+__all__ = [
+    'RandomForestRegressor',
+    'RandomForestClassifier',
+    'MultiTaskRandomForest',
+    'ExtraRandomTreesRegressor',
+    'ExtraRandomTreesClassifier'
+]
 
 
 class BaseLoloLearner(BaseEstimator, metaclass=ABCMeta):
@@ -86,19 +92,23 @@ def fit(self, X, y, weights=None, random_seed=None):
         assert weights_java.length() == len(X), "Weights copy failed"
 
         # Train the model
+        train_rows = self.gateway.jvm.io.citrine.lolo.api.TrainingRow.build(
+            train_data, self.gateway.jvm.scala.Some(weights_java)
+        )
         rng = self.gateway.jvm.io.citrine.lolo.util.LoloPyRandom.getRng(random_seed) if random_seed \
             else self.gateway.jvm.io.citrine.lolo.util.LoloPyRandom.getRng()
-        result = learner.train(train_data, self.gateway.jvm.scala.Some(weights_java), rng)
+        result = learner.train(train_rows, rng)
 
         # Unlink the training data, which is no longer needed (to save memory)
         self.gateway.detach(train_data)
         self.gateway.detach(weights_java)
+        self.gateway.detach(train_rows)
 
         # Get the model out
-        self.model_ = result.getModel()
+        self.model_ = result.model()
 
         # Store the feature importances
-        feature_importances_java = result.getFeatureImportance().get()
+        feature_importances_java = result.featureImportance().get()
         feature_importances_bytes = self.gateway.jvm.io.citrine.lolo.util.LoloPyDataLoader.send1DArray(feature_importances_java)
         self.feature_importances_ = np.frombuffer(feature_importances_bytes, 'float')
 
@@ -306,11 +316,8 @@ def predict_proba(self, X):
         return probs
 
 
-class RandomForestMixin(BaseLoloLearner):
-    """Random Forest base class
-
-    Implements the _make_learner operation and the __init__ function with options specific to the RandomForest
-    class in Lolo"""
+class RandomForestRegressor(BaseLoloRegressor):
+    """Random Forest model used for regression"""
 
     def __init__(self, num_trees=-1, use_jackknife=True, bias_learner=None,
                  leaf_learner=None, subset_strategy="auto", min_leaf_instances=1,
@@ -323,7 +330,7 @@ def __init__(self, num_trees=-1, use_jackknife=True, bias_learner=None,
             use_jackknife (bool): Whether to use jackknife based variance estimates
             bias_learner (BaseLoloLearner): Algorithm used to model bias (default: no model)
             leaf_learner (BaseLoloLearner): Learner used at each leaf of the random forest (default: GuessTheMean)
-            subset_strategy (Union[string,int,float]): Strategy used to determine number of features used at each split 
+            subset_strategy (Union[string,int,float]): Strategy used to determine number of features used at each split
                 Available options:
                     "auto": Use the default for lolo (all features for regression, sqrt for classification)
                     "log2": Use the base 2 log of the number of features
@@ -336,7 +343,7 @@ def __init__(self, num_trees=-1, use_jackknife=True, bias_learner=None,
             randomize_pivot_location (bool): whether to draw pivots randomly or always select the midpoint
             randomly_rotate_features (bool): whether to randomly rotate real features for each tree in the forest
         """
-        super(RandomForestMixin, self).__init__()
+        super().__init__()
 
         # Store the variables
         self.num_trees = num_trees
@@ -351,15 +358,21 @@ def __init__(self, num_trees=-1, use_jackknife=True, bias_learner=None,
         self.randomly_rotate_features = randomly_rotate_features
 
     def _make_learner(self):
-        #  TODO: Figure our a more succinct way of dealing with optional arguments/Option values
-        learner = self.gateway.jvm.io.citrine.lolo.learners.RandomForest(
-            self.num_trees, self.use_jackknife,
-            getattr(self.gateway.jvm.io.citrine.lolo.learners.RandomForest,
-                    "$lessinit$greater$default$3")() if self.bias_learner is None
-            else self.gateway.jvm.scala.Some(self.bias_learner._make_learner()),
-            getattr(self.gateway.jvm.io.citrine.lolo.learners.RandomForest,
-                    "$lessinit$greater$default$4")() if self.leaf_learner is None
-            else self.gateway.jvm.scala.Some(self.leaf_learner._make_learner()),
+        if self.bias_learner is None:
+            bias_learner = getattr(self.gateway.jvm.io.citrine.lolo.learners.RandomForestRegressor, "$lessinit$greater$default$3")()
+        else:
+            bias_learner = self.gateway.jvm.scala.Some(self.bias_learner._make_learner())
+
+        if self.leaf_learner is None:
+            leaf_learner = getattr(self.gateway.jvm.io.citrine.lolo.learners.RandomForestRegressor, "$lessinit$greater$default$4")()
+        else:
+            leaf_learner = self.gateway.jvm.scala.Some(self.leaf_learner._make_learner())
+
+        return self.gateway.jvm.io.citrine.lolo.learners.RandomForestRegressor(
+            self.num_trees,
+            self.use_jackknife,
+            bias_learner,
+            leaf_learner,
             self.subset_strategy,
             self.min_leaf_instances,
             self.max_depth,
@@ -367,22 +380,120 @@ def _make_learner(self):
             self.randomize_pivot_location,
             self.randomly_rotate_features
         )
-        return learner
 
 
-class RandomForestRegressor(BaseLoloRegressor, RandomForestMixin):
-    """Random Forest model used for regression"""
+class RandomForestClassifier(BaseLoloClassifier):
+    """Random Forest model used for classification"""
 
+    def __init__(self, num_trees=-1, use_jackknife=False, leaf_learner=None, subset_strategy="auto",
+                 min_leaf_instances=1, max_depth=2**30, randomize_pivot_location=False, randomly_rotate_features=False):
+        """Initialize the RandomForest
 
-class RandomForestClassifier(BaseLoloClassifier, RandomForestMixin):
-    """Random Forest model used for classification"""
+        Args:
+            num_trees (int): Number of trees to use in the forest (default of -1 sets the number of trees to the number of training rows)
+            use_jackknife (bool): Whether to use jackknife based variance estimates
+            leaf_learner (BaseLoloLearner): Learner used at each leaf of the random forest (default: GuessTheMean)
+            subset_strategy (Union[string,int,float]): Strategy used to determine number of features used at each split
+                Available options:
+                    "auto": Use the default for lolo (all features for regression, sqrt for classification)
+                    "log2": Use the base 2 log of the number of features
+                    "sqrt": Use the square root of the number of features
+                    integer: Set the number of features explicitly
+                    float: Use a certain fraction of the features
+            min_leaf_instances (int): Minimum number of features used at each leaf
+            max_depth (int): Maximum depth to which to allow the decision trees to grow
+            randomize_pivot_location (bool): whether to draw pivots randomly or always select the midpoint
+            randomly_rotate_features (bool): whether to randomly rotate real features for each tree in the forest
+        """
+        super().__init__()
 
+        # Store the variables
+        self.num_trees = num_trees
+        self.use_jackknife = use_jackknife
+        self.subset_strategy = subset_strategy
+        self.leaf_learner = leaf_learner
+        self.min_leaf_instances = min_leaf_instances
+        self.max_depth = max_depth
+        self.randomize_pivot_location = randomize_pivot_location
+        self.randomly_rotate_features = randomly_rotate_features
 
-class ExtraRandomTreesMixIn(BaseLoloLearner):
-    """Extra Random Trees base class
+    def _make_learner(self):
+        if self.leaf_learner is None:
+            leaf_learner = getattr(self.gateway.jvm.io.citrine.lolo.learners.RandomForestClassifier, "$lessinit$greater$default$3")()
+        else:
+            leaf_learner = self.gateway.jvm.scala.Some(self.leaf_learner._make_learner())
 
-    Implements the _make_learner operation and the __init__ function with options specific to the ExtraRandomTrees
-    class in Lolo"""
+        return self.gateway.jvm.io.citrine.lolo.learners.RandomForestClassifier(
+            self.num_trees,
+            self.use_jackknife,
+            leaf_learner,
+            self.subset_strategy,
+            self.min_leaf_instances,
+            self.max_depth,
+            self.randomize_pivot_location,
+            self.randomly_rotate_features
+        )
+
+
+class MultiTaskRandomForest(BaseLoloRegressor):
+    """Random Forest model used for regression on multiple outputs."""
+
+    def __init__(self, num_trees=-1, use_jackknife=True, bias_learner=None,
+                 subset_strategy="auto", min_leaf_instances=1, max_depth=2**30, uncertainty_calibration=False,
+                 randomize_pivot_location=False, randomly_rotate_features=False):
+        """Initialize the RandomForest
+
+        Args:
+            num_trees (int): Number of trees to use in the forest (default of -1 sets the number of trees to the number of training rows)
+            use_jackknife (bool): Whether to use jackknife based variance estimates
+            bias_learner (BaseLoloLearner): Algorithm used to model bias (default: no model)
+            subset_strategy (Union[string,int,float]): Strategy used to determine number of features used at each split
+                Available options:
+                    "auto": Use the default for lolo (all features for regression, sqrt for classification)
+                    "log2": Use the base 2 log of the number of features
+                    "sqrt": Use the square root of the number of features
+                    integer: Set the number of features explicitly
+                    float: Use a certain fraction of the features
+            min_leaf_instances (int): Minimum number of features used at each leaf
+            max_depth (int): Maximum depth to which to allow the decision trees to grow
+            uncertainty_calibration (bool): whether to re-calibrate the predicted uncertainty based on out-of-bag residuals
+            randomize_pivot_location (bool): whether to draw pivots randomly or always select the midpoint
+            randomly_rotate_features (bool): whether to randomly rotate real features for each tree in the forest
+        """
+        super().__init__()
+
+        # Store the variables
+        self.num_trees = num_trees
+        self.use_jackknife = use_jackknife
+        self.subset_strategy = subset_strategy
+        self.bias_learner = bias_learner
+        self.min_leaf_instances = min_leaf_instances
+        self.max_depth = max_depth
+        self.uncertainty_calibration = uncertainty_calibration
+        self.randomize_pivot_location = randomize_pivot_location
+        self.randomly_rotate_features = randomly_rotate_features
+
+    def _make_learner(self):
+        if self.bias_learner is None:
+            bias_learner = getattr(self.gateway.jvm.io.citrine.lolo.learners.MultiTaskRandomForest, "$lessinit$greater$default$3")()
+        else:
+            bias_learner = self.gateway.jvm.scala.Some(self.bias_learner._make_learner())
+
+        return self.gateway.jvm.io.citrine.lolo.learners.MultiTaskRandomForest(
+            self.num_trees,
+            self.use_jackknife,
+            bias_learner,
+            self.subset_strategy,
+            self.min_leaf_instances,
+            self.max_depth,
+            self.uncertainty_calibration,
+            self.randomize_pivot_location,
+            self.randomly_rotate_features
+        )
+
+
+class ExtraRandomTreesRegressor(BaseLoloRegressor):
+    """Extra Random Trees model used for regression."""
 
     def __init__(self, num_trees=-1, use_jackknife=False, bias_learner=None,
                  leaf_learner=None, subset_strategy="auto", min_leaf_instances=1,
@@ -405,15 +516,13 @@ def __init__(self, num_trees=-1, use_jackknife=False, bias_learner=None,
             min_leaf_instances (int): Minimum number of features used at each leaf
             max_depth (int): Maximum depth to which to allow the decision trees to grow
             uncertainty_calibration (bool): whether to re-calibrate the predicted uncertainty based on out-of-bag residuals
-            randomize_pivot_location (bool): whether to draw pivots randomly or always select the midpoint
-            disable_bootstrap (bool): whether to disable bootstrapping (default: true)
+            disable_bootstrap (bool): whether to disable bootstrapping (default: True)
             randomly_rotate_features (bool): whether to randomly rotate real features for each tree in the forest
         """
-        super(ExtraRandomTreesMixIn, self).__init__()
+        super().__init__()
 
         # Store the variables
         self.num_trees = num_trees
-
         self.use_jackknife = use_jackknife
         self.bias_learner = bias_learner
         self.leaf_learner = leaf_learner
@@ -421,18 +530,25 @@ def __init__(self, num_trees=-1, use_jackknife=False, bias_learner=None,
         self.min_leaf_instances = min_leaf_instances
         self.max_depth = max_depth
         self.uncertainty_calibration = uncertainty_calibration
+        self.disable_bootstrap = disable_bootstrap
         self.randomly_rotate_features = randomly_rotate_features
 
     def _make_learner(self):
-        #  TODO: Figure our a more succinct way of dealing with optional arguments/Option values
-        learner = self.gateway.jvm.io.citrine.lolo.learners.ExtraRandomTrees(
-            self.num_trees, self.use_jackknife,
-            getattr(self.gateway.jvm.io.citrine.lolo.learners.ExtraRandomTrees,
-                    "$lessinit$greater$default$3")() if self.bias_learner is None
-            else self.gateway.jvm.scala.Some(self.bias_learner._make_learner()),
-            getattr(self.gateway.jvm.io.citrine.lolo.learners.ExtraRandomTrees,
-                    "$lessinit$greater$default$4")() if self.leaf_learner is None
-            else self.gateway.jvm.scala.Some(self.leaf_learner._make_learner()),
+        if self.bias_learner is None:
+            bias_learner = getattr(self.gateway.jvm.io.citrine.lolo.learners.ExtraRandomTreesRegressor, "$lessinit$greater$default$3")()
+        else:
+            bias_learner = self.gateway.jvm.scala.Some(self.bias_learner._make_learner())
+
+        if self.leaf_learner is None:
+            leaf_learner = getattr(self.gateway.jvm.io.citrine.lolo.learners.ExtraRandomTreesRegressor, "$lessinit$greater$default$4")()
+        else:
+            leaf_learner = self.gateway.jvm.scala.Some(self.leaf_learner._make_learner())
+
+        return self.gateway.jvm.io.citrine.lolo.learners.ExtraRandomTreesRegressor(
+            self.num_trees,
+            self.use_jackknife,
+            bias_learner,
+            leaf_learner,
             self.subset_strategy,
             self.min_leaf_instances,
             self.max_depth,
@@ -440,15 +556,63 @@ def _make_learner(self):
             self.disable_bootstrap,
             self.randomly_rotate_features
         )
-        return learner
 
 
-class ExtraRandomTreesRegressor(BaseLoloRegressor, RandomForestMixin):
-    """Random Forest model used for regression"""
+class ExtraRandomTreesClassifier(BaseLoloClassifier):
+    """Extra Random Trees model used for classification."""
 
+    def __init__(self, num_trees=-1, use_jackknife=False, bias_learner=None,
+                 leaf_learner=None, subset_strategy="auto", min_leaf_instances=1,
+                 max_depth=2**30, uncertainty_calibration=False, disable_bootstrap=True,
+                 randomly_rotate_features=False):
+        """Initialize the ExtraRandomTrees ensemble
 
-class ExtraRandomTreesClassifier(BaseLoloClassifier, RandomForestMixin):
-    """Random Forest model used for regression"""
+        Args:
+            num_trees (int): Number of trees to use in the forest (default of -1 sets the number of trees to the number of training rows)
+            use_jackknife (bool): Whether to use jackknife based variance estimates (default: False)
+            leaf_learner (BaseLoloLearner): Learner used at each leaf of the random forest (default: GuessTheMean)
+            subset_strategy (Union[string,int,float]): Strategy used to determine number of features used at each split
+                Available options:
+                    "auto": Use the default for lolo (all features for regression; classification not supported)
+                    "log2": Use the base 2 log of the number of features
+                    "sqrt": Use the square root of the number of features
+                    integer: Set the number of features explicitly
+                    float: Use a certain fraction of the features
+            min_leaf_instances (int): Minimum number of features used at each leaf
+            max_depth (int): Maximum depth to which to allow the decision trees to grow
+            disable_bootstrap (bool): whether to disable bootstrapping (default: True)
+            randomly_rotate_features (bool): whether to randomly rotate real features for each tree in the forest
+        """
+        super().__init__()
+
+        # Store the variables
+        self.num_trees = num_trees
+        self.use_jackknife = use_jackknife
+        self.bias_learner = bias_learner
+        self.leaf_learner = leaf_learner
+        self.subset_strategy = subset_strategy
+        self.min_leaf_instances = min_leaf_instances
+        self.max_depth = max_depth
+        self.uncertainty_calibration = uncertainty_calibration
+        self.disable_bootstrap = disable_bootstrap
+        self.randomly_rotate_features = randomly_rotate_features
+
+    def _make_learner(self):
+        if self.leaf_learner is None:
+            leaf_learner = getattr(self.gateway.jvm.io.citrine.lolo.learners.ExtraRandomTreesClassifier, "$lessinit$greater$default$3")()
+        else:
+            leaf_learner = self.gateway.jvm.scala.Some(self.leaf_learner._make_learner())
+
+        return self.gateway.jvm.io.citrine.lolo.learners.ExtraRandomTreesClassifier(
+            self.num_trees,
+            self.use_jackknife,
+            leaf_learner,
+            self.subset_strategy,
+            self.min_leaf_instances,
+            self.max_depth,
+            self.disable_bootstrap,
+            self.randomly_rotate_features
+        )
 
 
 class RegressionTreeLearner(BaseLoloRegressor):
diff --git a/python/lolopy/version.py b/python/lolopy/version.py
index cd9057be..efec725d 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__ = "2.0.0"
+__version__ = "3.0.0"
diff --git a/src/main/scala/io/citrine/lolo/Learner.scala b/src/main/scala/io/citrine/lolo/Learner.scala
deleted file mode 100644
index fe2d35eb..00000000
--- a/src/main/scala/io/citrine/lolo/Learner.scala
+++ /dev/null
@@ -1,52 +0,0 @@
-package io.citrine.lolo
-
-import io.citrine.random.Random
-
-trait Learner extends Serializable {
-
-  /**
-    * Train a model
-    *
-    * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
-    * @param rng          random number generator for reproducibility
-    * @return training result containing a model
-    */
-  def train(
-      trainingData: Seq[(Vector[Any], Any)],
-      weights: Option[Seq[Double]] = None,
-      rng: Random = Random()
-  ): TrainingResult
-
-  /**
-    * Train a model with weights
-    *
-    * @param trainingData with weights in the form (features, label, weight)
-    * @param rng          random number generator for reproducibility
-    * @return training result containing a model
-    */
-  def train(trainingData: Seq[(Vector[Any], Any, Double)], rng: Random): TrainingResult = {
-    train(trainingData.map(r => (r._1, r._2)), Some(trainingData.map(_._3)), rng)
-  }
-
-}
-
-/**
-  * 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
-    * @param rng          random number generator for reproducibility
-    * @return A training result that encompasses model(s) for all labels.
-    */
-  def train(
-      trainingData: Seq[(Vector[Any], Vector[Any])],
-      weights: Option[Seq[Double]] = None,
-      rng: Random = Random()
-  ): MultiTaskTrainingResult
-}
diff --git a/src/main/scala/io/citrine/lolo/TrainingResult.scala b/src/main/scala/io/citrine/lolo/TrainingResult.scala
deleted file mode 100644
index 98138213..00000000
--- a/src/main/scala/io/citrine/lolo/TrainingResult.scala
+++ /dev/null
@@ -1,41 +0,0 @@
-package io.citrine.lolo
-
-// TODO (PLA-10433): Can we put a type parameter on TrainingResult?
-trait TrainingResult extends Serializable {
-
-  /**
-    * Get the model contained in the training result
-    *
-    * @return the model
-    */
-  def getModel(): Model[PredictionResult[Any]]
-
-  /**
-    * Get a measure of the importance of the model features
-    *
-    * @return feature influences as an array of doubles
-    */
-  def getFeatureImportance(): Option[Vector[Double]] = None
-
-  /**
-    * Get a measure of the loss of the model, e.g. RMS OOB error
-    *
-    * @return
-    */
-  def getLoss(): Option[Double] = None
-
-  /**
-    * Get the predicted vs actual values, e.g. from OOB
-    *
-    * @return seq of (feature vector, predicted value, and actual value)
-    */
-  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
-}
diff --git a/src/main/scala/io/citrine/lolo/api/Learner.scala b/src/main/scala/io/citrine/lolo/api/Learner.scala
new file mode 100644
index 00000000..bfca5541
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/api/Learner.scala
@@ -0,0 +1,26 @@
+package io.citrine.lolo.api
+
+import io.citrine.random.Random
+
+/**
+  * Base trait for a supervised learner that produces a model.
+  *
+  * @tparam T the label type of the data the learner is trained on
+  */
+trait Learner[T] extends Serializable {
+
+  /**
+    * Train a model on the provided training data.
+    *
+    * @param trainingData to train on
+    * @param rng          random number generator for reproducibility
+    * @return training result containing a model
+    */
+  def train(trainingData: Seq[TrainingRow[T]], rng: Random = Random()): TrainingResult[T]
+}
+
+/** A learner that trains on multiple labels, outputting a single model that makes predictions for all labels. */
+trait MultiTaskLearner extends Learner[Vector[Any]] {
+
+  def train(trainingData: Seq[TrainingRow[Vector[Any]]], rng: Random = Random()): MultiTaskTrainingResult
+}
diff --git a/src/main/scala/io/citrine/lolo/Model.scala b/src/main/scala/io/citrine/lolo/api/Model.scala
similarity index 58%
rename from src/main/scala/io/citrine/lolo/Model.scala
rename to src/main/scala/io/citrine/lolo/api/Model.scala
index bbf0fbc2..24b683bb 100644
--- a/src/main/scala/io/citrine/lolo/Model.scala
+++ b/src/main/scala/io/citrine/lolo/api/Model.scala
@@ -1,8 +1,8 @@
-package io.citrine.lolo
+package io.citrine.lolo.api
 
 import breeze.linalg.DenseMatrix
 
-trait Model[+T <: PredictionResult[Any]] extends Serializable {
+trait Model[+T] extends Serializable {
 
   /**
     * Apply the model to a seq of inputs
@@ -10,7 +10,7 @@ trait Model[+T <: PredictionResult[Any]] extends Serializable {
     * @param inputs to apply the model to
     * @return a PredictionResult which includes, at least, the expected outputs
     */
-  def transform(inputs: Seq[Vector[Any]]): T
+  def transform(inputs: Seq[Vector[Any]]): PredictionResult[T]
 
   /**
     * Compute Shapley feature attributions for a given input in this node's subtree
@@ -21,24 +21,23 @@ trait Model[+T <: PredictionResult[Any]] extends Serializable {
     *         One row per feature, each of length equal to the output dimension.
     *         The output dimension is 1 for single-task regression, or equal to the number of classification categories.
     */
-  def shapley(
-      input: Vector[Any],
-      omitFeatures: Set[Int] = Set()
-  ): Option[DenseMatrix[Double]] = None
-
+  def shapley(input: Vector[Any], omitFeatures: Set[Int] = Set()): Option[DenseMatrix[Double]] = None
 }
 
 /** A model that predicts a sequence of values, corresponding to multiple labels. */
-trait MultiTaskModel extends Model[MultiTaskModelPredictionResult] {
+trait MultiTaskModel extends Model[Vector[Any]] {
 
   /** The number of labels. Every prediction must have this length. */
-  val numLabels: Int
+  def numLabels: Int
 
   /** A boolean sequence indicating which labels are real-valued. Its length must be equal to `numLabels`. */
-  def getRealLabels: Seq[Boolean]
+  def realLabels: Seq[Boolean]
 
   /** Individual models corresponding to each label */
-  def getModels: Seq[Model[PredictionResult[Any]]]
+  def models: Seq[Model[Any]]
+
+  override def transform(inputs: Seq[Vector[Any]]): MultiTaskModelPredictionResult
+
 }
 
 /**
@@ -47,17 +46,11 @@ trait MultiTaskModel extends Model[MultiTaskModelPredictionResult] {
   * @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
+case class ParallelModels(models: Seq[Model[Any]], realLabels: Seq[Boolean]) extends MultiTaskModel {
+  override def numLabels: Int = models.length
 
-  override def transform(inputs: Seq[Vector[Any]]) =
-    new ParallelModelsPredictionResult(
-      models.map { model =>
-        model.transform(inputs).getExpected()
-      }.transpose
-    )
+  override def transform(inputs: Seq[Vector[Any]]): ParallelModelsPredictionResult = {
+    val predictions = models.map(_.transform(inputs).expected).toVector.transpose
+    ParallelModelsPredictionResult(predictions)
+  }
 }
diff --git a/src/main/scala/io/citrine/lolo/PredictionResult.scala b/src/main/scala/io/citrine/lolo/api/PredictionResult.scala
similarity index 76%
rename from src/main/scala/io/citrine/lolo/PredictionResult.scala
rename to src/main/scala/io/citrine/lolo/api/PredictionResult.scala
index 5e296890..345f9e13 100644
--- a/src/main/scala/io/citrine/lolo/PredictionResult.scala
+++ b/src/main/scala/io/citrine/lolo/api/PredictionResult.scala
@@ -1,8 +1,9 @@
-package io.citrine.lolo
+package io.citrine.lolo.api
 
 /**
   * Container for prediction results; must include expected values
-  * Created by maxhutch on 11/29/16.
+  *
+  * @tparam T the type of label variables being predicted
   */
 trait PredictionResult[+T] {
 
@@ -11,7 +12,7 @@ trait PredictionResult[+T] {
     *
     * @return expected value of each prediction
     */
-  def getExpected(): Seq[T]
+  def expected: Seq[T]
 
   /**
     * Get the "uncertainty" of the prediction
@@ -21,14 +22,14 @@ trait PredictionResult[+T] {
     * @param observational whether the uncertainty should account for observational uncertainty
     * @return uncertainty of each prediction
     */
-  def getUncertainty(observational: Boolean = true): Option[Seq[Any]] = None
+  def uncertainty(observational: Boolean = true): Option[Seq[Any]] = None
 
   /**
     * Get the training row scores for each prediction
     *
     * @return sequence (over predictions) of sequence (over training rows) of importances
     */
-  def getImportanceScores(): Option[Seq[Seq[Double]]] = None
+  def importanceScores: Option[Seq[Seq[Double]]] = None
 
   /**
     * Get the improvement (positive) or damage (negative) due to each training row on a prediction
@@ -36,14 +37,14 @@ trait PredictionResult[+T] {
     * @param actuals to assess the improvement or damage against
     * @return Sequence (over predictions) of sequence (over training rows) of influence
     */
-  def getInfluenceScores(actuals: Seq[Any]): Option[Seq[Seq[Double]]] = None
+  def influenceScores(actuals: Seq[Any]): Option[Seq[Seq[Double]]] = None
 
   /**
     * Get the gradient or sensitivity of each prediction
     *
     * @return a vector of doubles for each prediction
     */
-  def getGradient(): Option[Seq[Vector[Double]]] = None
+  def gradient: Option[Seq[Vector[Double]]] = None
 }
 
 /**
@@ -68,7 +69,7 @@ trait RegressionResult extends PredictionResult[Double] {
     * This statistic is related to the https://en.wikipedia.org/wiki/Prediction_interval
     * It does not include estimated bias, even if the regression result contains a bias estimate.
     */
-  def getStdDevObs(): Option[Seq[Double]] = None
+  def stdDevObs: Option[Seq[Double]] = None
 
   /**
     * Get the expected error of the observations, if possible
@@ -77,7 +78,7 @@ trait RegressionResult extends PredictionResult[Double] {
     * This statistic includes the contribution of the estimated bias.  E.g., for a normal distribution of predicted
     * means, the total error is sqrt(bias**2 + variance).
     */
-  def getTotalErrorObs(): Option[Seq[Double]] = None
+  def totalErrorObs: Option[Seq[Double]] = None
 
   /**
     * Get a quantile from the distribution of predicted observations, if possible
@@ -89,7 +90,7 @@ trait RegressionResult extends PredictionResult[Double] {
     *
     * @param quantile to get, taken between 0.0 and 1.0 (i.e. not a percentile)
     */
-  def getTotalErrorQuantileObs(quantile: Double): Option[Seq[Double]] = None
+  def totalErrorQuantileObs(quantile: Double): Option[Seq[Double]] = None
 
   /**
     * Get the expected error of the predicted mean observations, if possible
@@ -100,7 +101,7 @@ trait RegressionResult extends PredictionResult[Double] {
     * This statistic includes the contribution of the estimated bias.  E.g., for a normal distribution of predicted
     * means, the total error is sqrt(bias**2 + variance)
     */
-  def getTotalError(): Option[Seq[Double]] = None
+  def totalError: Option[Seq[Double]] = None
 
   /**
     * Get a quantile from the distribution of predicted means, if possible
@@ -109,7 +110,7 @@ trait RegressionResult extends PredictionResult[Double] {
     * corrected.
     * @param quantile to get, taken between 0.0 and 1.0 (i.e. not a percentile)
     */
-  def getTotalErrorQuantile(quantile: Double): Option[Seq[Double]] = None
+  def totalErrorQuantile(quantile: Double): Option[Seq[Double]] = None
 
   /**
     * Get the standard deviation of the distribution of predicted mean observations, if possible
@@ -119,7 +120,7 @@ trait RegressionResult extends PredictionResult[Double] {
     * This statistic is related to the variance in the bias-variance trade-off
     * https://en.wikipedia.org/wiki/Bias%E2%80%93variance_tradeoff
     */
-  def getStdDevMean(): Option[Seq[Double]] = None
+  def stdDevMean: Option[Seq[Double]] = None
 
   /**
     * Get a quantile from the distribution of predicted means, if possible
@@ -127,18 +128,15 @@ trait RegressionResult extends PredictionResult[Double] {
     * The distribution for which these quantiles are computed should have zero-mean (i.e. no bias)
     * @param quantile to get, taken between 0.0 and 1.0 (i.e. not a percentile)
     */
-  def getQuantileMean(quantile: Double): Option[Seq[Double]] = None
+  def quantileMean(quantile: Double): Option[Seq[Double]] = None
 
   /**
-    * **EXPERIMENTAL** Get the estimated bias of each prediction, if possible
+    * Get the estimated bias of each prediction, if possible.
     *
     * The bias is signed and can be subtracted from the prediction to improve accuracy.
     * See https://en.wikipedia.org/wiki/Bias%E2%80%93variance_tradeoff
-    *
-    * It is unclear if this method will be a stable member of the interface.  It should be reviewed
-    * before the next formal release.
     */
-  def getBias(): Option[Seq[Double]] = None
+  def bias: Option[Seq[Double]] = None
 
   /**
     * Get the "uncertainty", which is the TotalError if non-observational and the StdDevObs if observational
@@ -146,27 +144,27 @@ trait RegressionResult extends PredictionResult[Double] {
     * @param observational whether the uncertainty should account for observational uncertainty
     * @return uncertainty of each prediction
     */
-  override def getUncertainty(observational: Boolean = true): Option[Seq[Any]] = {
+  override def uncertainty(observational: Boolean = true): Option[Seq[Any]] = {
     if (observational) {
-      getTotalErrorObs()
+      totalErrorObs
     } else {
-      getTotalError()
+      totalError
     }
   }
 
-  def getQuantile(quantile: Double, observational: Boolean = true): Option[Seq[Double]] = {
+  def quantile(quantile: Double, observational: Boolean = true): Option[Seq[Double]] = {
     if (observational) {
-      getTotalErrorQuantileObs(quantile)
+      totalErrorQuantileObs(quantile)
     } else {
-      getTotalErrorQuantile(quantile)
+      totalErrorQuantile(quantile)
     }
   }
 }
 
 /** Container for predictions made on multiple labels simultaneously. */
-trait MultiTaskModelPredictionResult extends PredictionResult[Seq[Any]] {
+trait MultiTaskModelPredictionResult extends PredictionResult[Vector[Any]] {
 
-  override def getUncertainty(observational: Boolean = true): Option[Seq[Seq[Any]]] = None
+  override def uncertainty(observational: Boolean = true): Option[Seq[Seq[Any]]] = None
 
   /**
     * Get the correlation coefficients between the predictions made on two labels.
@@ -178,10 +176,8 @@ trait MultiTaskModelPredictionResult extends PredictionResult[Seq[Any]] {
     * @param observational whether the uncertainty correlation should take observational noise into account
     * @return optional sequence of correlation coefficients between specified labels for each prediction
     */
-  def getUncertaintyCorrelation(i: Int, j: Int, observational: Boolean = true): Option[Seq[Double]] = None
+  def uncertaintyCorrelation(i: Int, j: Int, observational: Boolean = true): Option[Seq[Double]] = None
 }
 
 /** 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
-}
+case class ParallelModelsPredictionResult(expected: Seq[Vector[Any]]) extends MultiTaskModelPredictionResult
diff --git a/src/main/scala/io/citrine/lolo/api/TrainingResult.scala b/src/main/scala/io/citrine/lolo/api/TrainingResult.scala
new file mode 100644
index 00000000..5812c995
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/api/TrainingResult.scala
@@ -0,0 +1,41 @@
+package io.citrine.lolo.api
+
+trait TrainingResult[+T] extends Serializable {
+
+  /**
+    * Get the model contained in the training result
+    *
+    * @return the model
+    */
+  def model: Model[T]
+
+  /**
+    * Get a measure of the importance of the model features
+    *
+    * @return feature influences as an array of doubles
+    */
+  def featureImportance: Option[Vector[Double]] = None
+
+  /**
+    * Get a measure of the loss of the model, e.g. RMS OOB error
+    *
+    * @return
+    */
+  def loss: Option[Double] = None
+
+  /**
+    * Get the predicted vs actual values, e.g. from OOB
+    *
+    * @return seq of (feature vector, predicted value, and actual value)
+    */
+  def predictedVsActual: Option[Seq[(Vector[Any], T, T)]] = None
+}
+
+trait MultiTaskTrainingResult extends TrainingResult[Vector[Any]] {
+
+  def models: Seq[Model[Any]]
+
+  override def model: MultiTaskModel
+
+  override def predictedVsActual: Option[Seq[(Vector[Any], Vector[Option[Any]], Vector[Option[Any]])]] = None
+}
diff --git a/src/main/scala/io/citrine/lolo/api/TrainingRow.scala b/src/main/scala/io/citrine/lolo/api/TrainingRow.scala
new file mode 100644
index 00000000..63cd2b22
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/api/TrainingRow.scala
@@ -0,0 +1,49 @@
+package io.citrine.lolo.api
+
+/**
+  * A bundle of (inputs, label, weight) used for training a [[Learner]].
+  *
+  * @param inputs vector of mixed-type input features
+  * @param label  target variable to train a model on
+  * @param weight of the training row
+  * @tparam T the type of the label data
+  */
+case class TrainingRow[+T](inputs: Vector[Any], label: T, weight: Double = 1.0) {
+
+  /** Cast the training row as a tuple of (inputs, label, weight). */
+  def asTuple: (Vector[Any], T, Double) = (inputs, label, weight)
+
+  /** Return a new row updated with the provided inputs. */
+  def withInputs(inputs: Vector[Any]): TrainingRow[T] = copy(inputs = inputs)
+
+  /** Return a new row updated with the provided label. */
+  def withLabel[T2](label: T2): TrainingRow[T2] = copy(label = label)
+
+  /** Return a new row updated with the provided weight. */
+  def withWeight(weight: Double): TrainingRow[T] = copy(weight = weight)
+
+  /** Return a new row with the function applied to the original inputs. */
+  def mapInputs(f: Vector[Any] => Vector[Any]): TrainingRow[T] = copy(inputs = f(inputs))
+
+  /** Return a new row with the function applied to the original label. */
+  def mapLabel[T2](f: T => T2): TrainingRow[T2] = copy(label = f(label))
+
+  /** Return a new row with the function applied to the original weight. */
+  def mapWeight(f: Double => Double): TrainingRow[T] = copy(weight = f(weight))
+}
+
+object TrainingRow {
+
+  /** Build a sequence of training from the provided data and weights. */
+  def build[T](data: Seq[(Vector[Any], T)], weights: Option[Seq[Double]] = None): Seq[TrainingRow[T]] = {
+    val actualWeights = weights.getOrElse(Seq.fill(data.length)(1.0))
+    data.zip(actualWeights).map {
+      case ((i, l), w) => TrainingRow(i, l, w)
+    }
+  }
+
+  /** Build a sequence of training from the provided data and weights. */
+  def build[T](data: Seq[(Vector[Any], T, Double)]): Seq[TrainingRow[T]] = {
+    data.map { case (i, l, w) => TrainingRow(i, l, w) }
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/bags/BaggedModel.scala b/src/main/scala/io/citrine/lolo/bags/BaggedModel.scala
new file mode 100644
index 00000000..0e6171ca
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/bags/BaggedModel.scala
@@ -0,0 +1,75 @@
+package io.citrine.lolo.bags
+
+import breeze.linalg.DenseMatrix
+import io.citrine.lolo.api.Model
+
+import scala.collection.parallel.immutable.ParSeq
+
+/** A model holding a parallel sequence of models and the sample counts used to train them. */
+trait BaggedModel[+T] extends Model[T] {
+
+  /** Models in the ensemble trained on subsets of the training data. */
+  def ensembleModels: ParSeq[Model[T]]
+
+  override def transform(inputs: Seq[Vector[Any]]): BaggedPrediction[T]
+
+  override def shapley(input: Vector[Any], omitFeatures: Set[Int] = Set()): Option[DenseMatrix[Double]] = {
+    val ensembleShapley = ensembleModels.map(model => model.shapley(input, omitFeatures))
+    if (ensembleShapley.head.isEmpty) {
+      None
+    } else {
+      assert(ensembleShapley.forall(x => x.isDefined))
+
+      def sumReducer(a: Option[DenseMatrix[Double]], b: Option[DenseMatrix[Double]]): Option[DenseMatrix[Double]] = {
+        Some(a.get + b.get)
+      }
+
+      val scale = 1.0 / ensembleShapley.length
+      Some(scale * ensembleShapley.reduce(sumReducer).get)
+    }
+  }
+}
+
+case class BaggedRegressionModel(
+    ensembleModels: ParSeq[Model[Double]],
+    Nib: Vector[Vector[Int]],
+    rescaleRatio: Double = 1.0,
+    disableBootstrap: Boolean = false,
+    biasModel: Option[Model[Double]] = None
+) extends BaggedModel[Double] {
+
+  override def transform(inputs: Seq[Vector[Any]]): BaggedRegressionPrediction = {
+    assert(inputs.forall(_.size == inputs.head.size))
+
+    val bias = biasModel.map(_.transform(inputs).expected)
+    val ensemblePredictions = ensembleModels.map(model => model.transform(inputs)).seq
+
+    if (inputs.size == 1) {
+      // In the special case of a single prediction on a real value, emit an optimized prediction class
+      SinglePointBaggedPrediction(
+        ensemblePredictions,
+        Nib,
+        bias.map(_.head),
+        rescaleRatio,
+        disableBootstrap
+      )
+    } else {
+      MultiPointBaggedPrediction(
+        ensemblePredictions,
+        Nib,
+        bias,
+        rescaleRatio,
+        disableBootstrap
+      )
+    }
+  }
+}
+
+case class BaggedClassificationModel[T](ensembleModels: ParSeq[Model[T]]) extends BaggedModel[T] {
+
+  override def transform(inputs: Seq[Vector[Any]]): BaggedClassificationPrediction[T] = {
+    assert(inputs.forall(_.size == inputs.head.size))
+    val ensemblePredictions = ensembleModels.map(model => model.transform(inputs)).seq
+    BaggedClassificationPrediction(ensemblePredictions)
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/bags/BaggedResult.scala b/src/main/scala/io/citrine/lolo/bags/BaggedPrediction.scala
similarity index 58%
rename from src/main/scala/io/citrine/lolo/bags/BaggedResult.scala
rename to src/main/scala/io/citrine/lolo/bags/BaggedPrediction.scala
index 5c3746f5..dae7a9a8 100644
--- a/src/main/scala/io/citrine/lolo/bags/BaggedResult.scala
+++ b/src/main/scala/io/citrine/lolo/bags/BaggedPrediction.scala
@@ -1,134 +1,119 @@
 package io.citrine.lolo.bags
 
 import breeze.linalg.{norm, DenseMatrix, DenseVector}
-import io.citrine.lolo.stats.StatsUtils
-import io.citrine.lolo.{
-  MultiTaskModelPredictionResult,
-  ParallelModelsPredictionResult,
-  PredictionResult,
-  RegressionResult
-}
+import io.citrine.lolo.api.{PredictionResult, RegressionResult}
 import org.slf4j.{Logger, LoggerFactory}
 
 /**
-  * Interface defining the return value of a [[BaggedModel]]
+  * Interface defining the return value of a [[BaggedModel]].
+  *
+  * This allows the implementation to depend on the number of simultaneous predictions,
+  * which has performance implications.
   *
-  * 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]]
+trait BaggedPrediction[+T] extends PredictionResult[T] {
 
-  /** The number of inputs that have been predicted on (NOT the number of bagged models). */
+  /** The number of input rows that have been predicted on (NOT the number of bagged models). */
   def numPredictions: Int
 
+  /** The predictions made by each of the bagged models in the ensemble. */
+  def ensemblePredictions: Seq[PredictionResult[T]]
+
   /**
     * Average the gradients from the models in the ensemble
     *
     * @return the gradient of each prediction as a vector of doubles
     */
-  override def getGradient(): Option[Seq[Vector[Double]]] = gradient
-
-  private lazy val gradient = if (predictions.head.getGradient().isEmpty) {
+  override lazy val gradient: Option[Seq[Vector[Double]]] = if (ensemblePredictions.head.gradient.isEmpty) {
     /* If the underlying model has no gradient, return None */
     None
   } else {
-    val gradientsByPrediction: Seq[Seq[Vector[Double]]] = predictions.map(_.getGradient().get)
+    val gradientsByPrediction: Seq[Seq[Vector[Double]]] = ensemblePredictions.map(_.gradient.get)
     val gradientsByInput: Seq[Seq[Vector[Double]]] = gradientsByPrediction.transpose
-    Some(gradientsByInput.map { r =>
-      r.toVector.transpose.map(_.sum / predictions.size)
-    })
+    Some(
+      gradientsByInput.map { r => r.toVector.transpose.map(_.sum / ensemblePredictions.size) }
+    )
   }
 }
 
+/** Interface for [[BaggedPrediction]]s for regression tasks with task-specific methods. */
+sealed trait BaggedRegressionPrediction extends BaggedPrediction[Double] with RegressionResult
+
 /**
   * 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.
+  * See [[MultiPointBaggedPrediction]] 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
+  * @param ensemblePredictions for each constituent model
+  * @param NibIn               the sample matrix as (N_models x N_training)
+  * @param pointBias           bias model prediction at the single prediction point
   */
-case class SinglePredictionBaggedResult(
-    predictions: Seq[PredictionResult[Double]],
+case class SinglePointBaggedPrediction(
+    ensemblePredictions: Seq[PredictionResult[Double]],
     NibIn: Vector[Vector[Int]],
-    bias: Option[Double] = None,
-    rescale: Double = 1.0,
+    pointBias: Option[Double] = None,
+    rescaleRatio: Double = 1.0,
     disableBootstrap: Boolean = false
-) extends BaggedResult[Double]
-    with RegressionResult {
-  private lazy val treePredictions: Array[Double] = predictions.map(_.getExpected().head).toArray
+) extends BaggedRegressionPrediction {
 
   override def numPredictions: Int = 1
 
+  override def expected: Seq[Double] = Seq(rawExpected + pointBias.getOrElse(0.0))
+
+  override def bias: Option[Seq[Double]] = pointBias.map(Seq(_))
+
+  override def uncertainty(observational: Boolean): Option[Seq[Double]] = {
+    if (observational) {
+      stdDevObs
+    } else {
+      stdDevMean
+    }
+  }
+
   /**
-    * Return the ensemble average or maximum vote
+    * The importances are computed as an average of bias-corrected jackknife-after-bootstrap
+    * and infinitesimal jackknife methods
     *
-    * @return expected value of each prediction
+    * @return training row scores of each prediction
     */
-  override def getExpected(): Seq[Double] = Seq(expected + bias.getOrElse(0.0))
+  override def importanceScores: Option[Seq[Seq[Double]]] = Some(Seq(singleScores))
 
-  private lazy val expected = treePredictions.sum / treePredictions.length
-  private lazy val treeVariance: Double = {
-    assert(treePredictions.length > 1, "Bootstrap variance undefined for fewer than 2 bootstrap samples.")
-    treePredictions.map(x => Math.pow(x - expected, 2.0)).sum / (treePredictions.length - 1)
-  }
-
-  override def getStdDevMean(): Option[Seq[Double]] = {
+  override lazy val stdDevMean: Option[Seq[Double]] = {
     if (disableBootstrap) {
       // If bootstrap is disabled, rescale is unity and treeVariance is our only option for UQ.
       // Since it's not recalibrated, it's best considered to be a confidence interval of the underlying weak learner.
-      assert(rescale == 1.0)
-      Some(Seq(Math.sqrt(treeVariance)))
+      assert(rescaleRatio == 1.0)
+      Some(Seq(math.sqrt(treeVariance)))
     } else {
-      Some(Seq(stdDevMean))
+      Some(
+        Seq(math.sqrt(BaggedPrediction.rectifyEstimatedVariance(singleScores)))
+      )
     }
   }
 
-  override def getStdDevObs(): Option[Seq[Double]] = {
-    if (disableBootstrap) {
-      None
-    } else {
-      Some(Seq(stdDevObs))
-    }
+  override lazy val stdDevObs: Option[Seq[Double]] = Option.when(!disableBootstrap) {
+    val std = (rescaleRatio * math.sqrt(treeVariance)) ensuring (_ >= 0.0)
+    Seq(std)
   }
 
-  /**
-    * For the sake of parity, we were using this method
-    */
-  override def getUncertainty(observational: Boolean): Option[Seq[Any]] = {
-    if (observational) {
-      getStdDevObs()
-    } else {
-      getStdDevMean()
-    }
+  private lazy val treePredictions = ensemblePredictions.map(_.expected.head).toArray
+  private lazy val rawExpected = treePredictions.sum / treePredictions.length
+  private lazy val treeVariance = {
+    assert(treePredictions.length > 1, "Bootstrap variance undefined for fewer than 2 bootstrap samples.")
+    treePredictions.map(x => math.pow(x - rawExpected, 2.0)).sum / (treePredictions.length - 1)
   }
 
-  private lazy val stdDevMean: Double = Math.sqrt(BaggedResult.rectifyEstimatedVariance(singleScores))
-
-  private lazy val stdDevObs: Double = {
-    rescale * Math.sqrt(treeVariance)
-  } ensuring (_ >= 0.0)
-
-  /**
-    * The importances are computed as an average of bias-corrected jackknife-after-bootstrap
-    * and infinitesimal jackknife methods
-    *
-    * @return training row scores of each prediction
-    */
-  override def getImportanceScores(): Option[Seq[Seq[Double]]] = Some(Seq(singleScores))
-
   private lazy val singleScores: Vector[Double] = {
     // 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)
+    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 =>
@@ -142,7 +127,7 @@ case class SinglePredictionBaggedResult(
       var tNot: Double = 0.0
       var tNotCount: Int = 0
       vecN.indices.foreach { jdx =>
-        cov = cov + (vecN(jdx) - 1) * (treePredictions(jdx) - expected)
+        cov = cov + (vecN(jdx) - 1) * (treePredictions(jdx) - rawExpected)
 
         if (vecN(jdx) == 0) {
           tNot = tNot + treePredictions(jdx)
@@ -150,13 +135,13 @@ case class SinglePredictionBaggedResult(
         }
       }
       // Compute the infinitesimal jackknife estimate
-      val varIJ = Math.pow(cov / vecN.size, 2.0)
+      val varIJ = math.pow(cov / vecN.length, 2.0)
 
       if (tNotCount > 0) {
         // Compute the Jackknife after bootstrap estimate
-        val varJ = Math.pow(tNot / tNotCount - expected, 2.0) * (nMat.size - 1) / nMat.size
+        val varJ = math.pow(tNot / tNotCount - rawExpected, 2.0) * (nMat.size - 1) / nMat.size
         // Averaged the corrected IJ and J estimates
-        0.5 * (varJ + varIJ - Math.E * correction)
+        0.5 * (varJ + varIJ - math.E * correction)
       } else {
         // We can't compute the Jackknife after bootstrap estimate, so just correct the IJ estimate
         varIJ - correction
@@ -168,137 +153,91 @@ case class SinglePredictionBaggedResult(
   }
 }
 
-case class BaggedClassificationResult(
-    predictions: Seq[PredictionResult[Any]]
-) extends BaggedResult[Any] {
-  val predictionEnsemble = predictions.map { p => p.getExpected() }
-  lazy val expectedMatrix: Seq[Seq[Any]] = predictions.map(p => p.getExpected()).transpose
-
-  lazy val expected: Seq[Any] = expectedMatrix.map(ps => ps.groupBy(identity).maxBy(_._2.size)._1)
-  lazy val uncertainty: Seq[Map[Any, Double]] =
-    expectedMatrix.map(ps => ps.groupBy(identity).view.mapValues(_.size.toDouble / ps.size).toMap)
-
-  override def numPredictions: Int = expectedMatrix.length
-
-  /**
-    * Return the majority vote vote
-    *
-    * @return expected value of each prediction
-    */
-  override def getExpected(): Seq[Any] = expected
-
-  override def getUncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = Some(uncertainty)
-}
-
 /**
   * Container with model-wise predictions and logic to compute variances and training row scores
   *
   * 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 [[SinglePredictionBaggedResult]], which is more clear.  These two implementations are tested for consistency.
+  * look at [[SinglePointBaggedPrediction]], which is more clear.  These two implementations are tested for consistency.
   *
-  * @param predictions for each constituent model
+  * @param ensemblePredictions for each constituent model
   * @param NibIn       the sample matrix as (N_models x N_training)
-  * @param bias        model to use for estimating bias
+  * @param bias        bias model predictions at all of the prediction points
   */
-case class MultiPredictionBaggedResult(
-    predictions: Seq[PredictionResult[Double]],
+case class MultiPointBaggedPrediction(
+    ensemblePredictions: Seq[PredictionResult[Double]],
     NibIn: Vector[Vector[Int]],
-    bias: Option[Seq[Double]] = None,
-    rescale: Double = 1.0,
+    override val bias: Option[Seq[Double]] = None,
+    rescaleRatio: Double = 1.0,
     disableBootstrap: Boolean = false
-) extends BaggedResult[Double]
-    with RegressionResult {
+) extends BaggedRegressionPrediction {
 
-  /**
-    * Return the ensemble average
-    *
-    * @return expected value of each prediction
-    */
-  override def getExpected(): Seq[Double] = expected.zip(biasCorrection).map(x => x._1 + x._2)
-
-  override def getStdDevObs(): Option[Seq[Double]] = {
-    if (disableBootstrap) {
-      None
-    } else {
-      Some(varObs.map { v => Math.sqrt(v) })
-    }
-  }
+  override def numPredictions: Int = expectedMatrix.length
 
-  override def getStdDevMean(): Option[Seq[Double]] = {
-    if (disableBootstrap) {
-      // If bootstrap is disabled, rescale is unity and treeVariance is our only option for UQ.
-      // Since it's not recalibrated, it's best considered to be a confidence interval of the underlying weak learner.
-      assert(rescale == 1.0)
-      Some(varObs.map { v => Math.sqrt(v) })
-    } else {
-      Some(stdDevMean)
-    }
-  }
+  override def expected: Seq[Double] = rawExpected.zip(biasCorrection).map(x => x._1 + x._2)
 
-  /**
-    * For the sake of parity, we were using this method
-    */
-  override def getUncertainty(observational: Boolean): Option[Seq[Any]] = {
+  override def uncertainty(observational: Boolean): Option[Seq[Any]] = {
     if (observational) {
-      getStdDevObs()
+      stdDevObs
     } else {
-      getStdDevMean()
+      stdDevMean
     }
   }
 
-  /**
-    * Return IJ scores
-    *
-    * @return training row scores of each prediction
-    */
-  override def getInfluenceScores(actuals: Seq[Any]): Option[Seq[Seq[Double]]] = {
+  override def influenceScores(actuals: Seq[Any]): Option[Seq[Seq[Double]]] = {
     Some(
       influences(
-        expected.asInstanceOf[Seq[Double]].toVector,
+        expected.toVector,
         actuals.toVector.asInstanceOf[Vector[Double]],
-        expectedMatrix.asInstanceOf[Seq[Seq[Double]]],
+        expectedMatrix,
         NibJMat,
         NibIJMat
       )
     )
   }
 
-  override def getImportanceScores(): Option[Seq[Seq[Double]]] = Some(scores)
+  override def importanceScores: Option[Seq[Seq[Double]]] = Some(scores)
 
-  override def numPredictions: Int = expectedMatrix.length
+  override lazy val stdDevObs: Option[Seq[Double]] = Option.when(!disableBootstrap)(varObs.map(math.sqrt))
+
+  override lazy val stdDevMean: Option[Seq[Double]] = Option
+    .when(disableBootstrap) {
+      // If bootstrap is disabled, rescale is unity and treeVariance is our only option for UQ.
+      // Since it's not recalibrated, it's best considered to be a confidence interval of the underlying weak learner.
+      assert(rescaleRatio == 1.0)
+      varObs.map(math.sqrt)
+    }
+    .orElse {
+      val std = variance(rawExpected.toVector, expectedMatrix, NibJMat, NibIJMat).map(math.sqrt)
+      Some(std)
+    }
 
   /* transpose to be (# training) x (# models) */
-  lazy val Nib: Vector[Vector[Int]] = NibIn.transpose
+  private 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
+  private lazy val expectedMatrix: Seq[Seq[Double]] = ensemblePredictions.map(p => p.expected).transpose
 
   /* 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)
+  private lazy val biasCorrection: Seq[Double] = bias.getOrElse(Seq.fill(expectedMatrix.length)(0))
+  private lazy val rawExpected: Seq[Double] = expectedMatrix.map(ps => ps.sum / ps.size)
 
-  lazy val NibJMat = BaggedResult.getJackknifeAfterBootstrapMatrix(Nib)
+  private lazy val NibJMat: DenseMatrix[Double] = BaggedPrediction.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 }
+  private lazy val NibIJMat: DenseMatrix[Double] = BaggedPrediction.getInfinitesimalJackknifeMatrix(Nib)
 
   /* This estimates the variance of predictive distribution. */
-  lazy val varObs: Seq[Double] =
-    expectedMatrix.asInstanceOf[Seq[Seq[Double]]].zip(expected.asInstanceOf[Seq[Double]]).map {
-      case (b, y) =>
-        assert(Nib.size > 1, "Bootstrap variance undefined for fewer than 2 bootstrap samples.")
-        b.map { x => rescale * rescale * Math.pow(x - y, 2.0) }.sum / (b.size - 1)
-    }
+  private lazy val varObs: Seq[Double] = expectedMatrix.zip(rawExpected).map {
+    case (b, y) =>
+      assert(Nib.size > 1, "Bootstrap variance undefined for fewer than 2 bootstrap samples.")
+      b.map { x => rescaleRatio * rescaleRatio * math.pow(x - y, 2.0) }.sum / (b.size - 1)
+  }
 
   /* Compute the scores one prediction at a time */
-  lazy val scores: Seq[Vector[Double]] = scores(expected.toVector, expectedMatrix, NibJMat, NibIJMat)
+  private lazy val scores: Seq[Vector[Double]] = scores(expected.toVector, expectedMatrix, NibJMat, NibIJMat)
     // make sure the variance is non-negative after the stochastic correction
-    .map(BaggedResult.rectifyImportanceScores)
-    .map(_.map(Math.sqrt))
+    .map(BaggedPrediction.rectifyImportanceScores)
+    .map(_.map(math.sqrt))
 
   /**
     * Compute the variance of a prediction as the average of bias corrected IJ and J variance estimates
@@ -315,7 +254,7 @@ case class MultiPredictionBaggedResult(
       NibJ: DenseMatrix[Double],
       NibIJ: DenseMatrix[Double]
   ): Seq[Double] = {
-    scores(meanPrediction, modelPredictions, NibJ, NibIJ).map { BaggedResult.rectifyEstimatedVariance }
+    scores(meanPrediction, modelPredictions, NibJ, NibIJ).map { BaggedPrediction.rectifyEstimatedVariance }
   }
 
   /**
@@ -345,15 +284,15 @@ case class MultiPredictionBaggedResult(
     val arg = IJMat2 + JMat2
 
     /* 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
+    val prefactor = 1.0 / math.pow(modelPredictions.head.size, 2.0) * (Nib.size - 1) / Nib.size
 
     modelPredictions.indices
       .map { i =>
         /* Compute the first order bias correction for the variance estimators */
-        val correction = prefactor * Math.pow(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)
+        0.5 * (arg(::, i) - math.E * correction)
       }
       .map(_.toScalaVector)
   }
@@ -393,10 +332,10 @@ case class MultiPredictionBaggedResult(
 
         /* The correction is prediction dependent, so we need to operate on vectors */
         val influencePerRow: DenseVector[Double] =
-          Math.signum(actualPrediction(i) - meanPrediction(i)) * 0.5 * (arg(::, i) - Math.E * correction)
+          math.signum(actualPrediction(i) - meanPrediction(i)) * 0.5 * (arg(::, i) - math.E * correction)
 
         /* Impose a floor in case any of the variances are negative (hacked to work in breeze) */
-        // val floor: Double = Math.min(0, -min(variancePerRow))
+        // val floor: Double = math.min(0, -min(variancePerRow))
         // val rezero: DenseVector[Double] = variancePerRow - floor
         // 0.5 * (rezero + abs(rezero)) + floor
         influencePerRow
@@ -405,78 +344,22 @@ case class MultiPredictionBaggedResult(
   }
 }
 
-/**
-  * 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))
+case class BaggedClassificationPrediction[T](ensemblePredictions: Seq[PredictionResult[T]])
+    extends BaggedPrediction[T] {
 
-  // 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 numPredictions: Int = expectedMatrix.length
 
-  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
-    }
-  }
+  override def uncertainty(includeNoise: Boolean = true): Option[Seq[Map[T, Double]]] = Some(predictionUncertainties)
 
-  /** 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)
-    }
-  }
+  override lazy val expected: Seq[T] = expectedMatrix.map(ps => ps.groupBy(identity).maxBy(_._2.size)._1)
 
-  /**
-    * 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)
+  private lazy val expectedMatrix: Seq[Seq[T]] = ensemblePredictions.map(p => p.expected).transpose
 
+  private lazy val predictionUncertainties: Seq[Map[T, Double]] =
+    expectedMatrix.map(ps => ps.groupBy(identity).view.mapValues(_.size.toDouble / ps.size).toMap)
 }
 
-object BaggedResult {
+object BaggedPrediction {
 
   val logger: Logger = LoggerFactory.getLogger(getClass)
 
@@ -560,7 +443,7 @@ object BaggedResult {
       logger.warn(
         s"All scores were negative; using the magnitude of the smallest score as an estimate for the variance.  It is highly recommended to increase the ensemble size."
       )
-      -scores.min // equivalent to Math.abs(scores.min)
+      -scores.min // equivalent to math.abs(scores.min)
     }
   } ensuring (_ >= 0.0)
 
@@ -579,11 +462,10 @@ object BaggedResult {
     */
   def rectifyImportanceScores(scores: Vector[Double]): Vector[Double] = {
     // this is a lower-bound on the noise level; note that it is strictly smaller than the correction
-    val floor = Math.abs(scores.min)
-
+    val floor = math.abs(scores.min)
     if (floor < 0.0) {
       logger.warn(s"Some importance scores were negative; rectifying.  Please consider increasing the ensemble size.")
     }
-    scores.map(Math.max(floor, _))
+    scores.map(math.max(floor, _))
   } ensuring (vec => vec.forall(_ >= 0.0))
 }
diff --git a/src/main/scala/io/citrine/lolo/bags/BaggedTrainingResult.scala b/src/main/scala/io/citrine/lolo/bags/BaggedTrainingResult.scala
new file mode 100644
index 00000000..fa27ab40
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/bags/BaggedTrainingResult.scala
@@ -0,0 +1,85 @@
+package io.citrine.lolo.bags
+
+import io.citrine.lolo.api.{Model, TrainingResult, TrainingRow}
+import io.citrine.lolo.stats.metrics.{ClassificationMetrics, RegressionMetrics}
+
+import scala.collection.parallel.immutable.ParSeq
+
+/** The result of training a [[Bagger]] to produce a [[BaggedModel]]. */
+sealed trait BaggedTrainingResult[+T] extends TrainingResult[T] {
+
+  override def model: BaggedModel[T]
+}
+
+case class RegressionBaggerTrainingResult(
+    ensembleModels: ParSeq[Model[Double]],
+    Nib: Vector[Vector[Int]],
+    trainingData: Seq[TrainingRow[Double]],
+    override val featureImportance: Option[Vector[Double]],
+    biasModel: Option[Model[Double]] = None,
+    rescaleRatio: Double = 1.0,
+    disableBootstrap: Boolean = false
+) extends BaggedTrainingResult[Double] {
+
+  lazy val NibT: Vector[Vector[Int]] = Nib.transpose
+
+  override lazy val model: BaggedRegressionModel =
+    BaggedRegressionModel(ensembleModels, Nib, rescaleRatio, disableBootstrap, biasModel)
+
+  override lazy val predictedVsActual: Option[Seq[(Vector[Any], Double, Double)]] = Some(
+    trainingData.zip(NibT).flatMap {
+      case (TrainingRow(f, l, _), nb) =>
+        val oob = if (disableBootstrap) {
+          ensembleModels.zip(nb)
+        } else {
+          ensembleModels.zip(nb).filter(_._2 == 0)
+        }
+
+        if (oob.isEmpty || l.isNaN) {
+          Seq.empty
+        } else {
+          val predicted = oob.map(_._1.transform(Seq(f)).expected.head).sum / oob.size
+          Seq((f, predicted, l))
+        }
+    }
+  )
+
+  override lazy val loss: Option[Double] = predictedVsActual.collect {
+    case pva if pva.nonEmpty => RegressionMetrics.RMSE(pva)
+  }
+}
+
+case class ClassificationBaggerTrainingResult[T](
+    ensembleModels: ParSeq[Model[T]],
+    Nib: Vector[Vector[Int]],
+    trainingData: Seq[TrainingRow[T]],
+    override val featureImportance: Option[Vector[Double]],
+    disableBootstrap: Boolean = false
+) extends BaggedTrainingResult[T] {
+
+  lazy val NibT: Vector[Vector[Int]] = Nib.transpose
+
+  override lazy val model: BaggedClassificationModel[T] = BaggedClassificationModel(ensembleModels)
+
+  override lazy val predictedVsActual: Option[Seq[(Vector[Any], T, T)]] = Some(
+    trainingData.zip(NibT).flatMap {
+      case (TrainingRow(f, l, _), nb) =>
+        val oob = if (disableBootstrap) {
+          ensembleModels.zip(nb)
+        } else {
+          ensembleModels.zip(nb).filter(_._2 == 0)
+        }
+
+        if (oob.isEmpty || l == null) {
+          Seq()
+        } else {
+          val predicted = oob.map(_._1.transform(Seq(f)).expected.head).groupBy(identity).maxBy(_._2.size)._1
+          Seq((f, predicted, l))
+        }
+    }
+  )
+
+  override lazy val loss: Option[Double] = predictedVsActual.collect {
+    case pva if pva.nonEmpty => ClassificationMetrics.loss(pva)
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/bags/Bagger.scala b/src/main/scala/io/citrine/lolo/bags/Bagger.scala
index 14b7349e..d60d7e4f 100644
--- a/src/main/scala/io/citrine/lolo/bags/Bagger.scala
+++ b/src/main/scala/io/citrine/lolo/bags/Bagger.scala
@@ -1,84 +1,89 @@
 package io.citrine.lolo.bags
 
-import breeze.linalg.DenseMatrix
 import breeze.stats.distributions.Poisson
+import io.citrine.lolo.api.{Learner, Model, TrainingRow}
+import io.citrine.lolo.bags.Bagger.BaggedEnsemble
+import io.citrine.lolo.stats.StatsUtils
 import io.citrine.random.Random
-import io.citrine.lolo.stats.metrics.{ClassificationMetrics, RegressionMetrics}
-import io.citrine.lolo.stats.StatsUtils.breezeRandBasis
-import io.citrine.lolo.{Learner, Model, PredictionResult, TrainingResult}
 
-import scala.collection.parallel.immutable.ParSeq
-import scala.collection.parallel.CollectionConverters.IterableIsParallelizable
-import scala.reflect._
+import scala.collection.parallel.CollectionConverters._
+import scala.collection.parallel.immutable.ParVector
 
-/**
-  * A bagger creates an ensemble of models by training the learner on random samples of the training data
-  *
-  * @param method  learner to train each model in the ensemble
-  * @param numBags number of base models to aggregate (default of -1 sets the number of models to the number of training rows)
-  * @param useJackknife whether to enable jackknife uncertainty estimate
-  * @param biasLearner learner to use for estimating bias
-  * @param uncertaintyCalibration whether to enable empirical uncertainty calibration
-  * @param disableBootstrap whether to disable bootstrap (useful when `method` implements its own randomization)
-  */
-case class Bagger(
-    method: Learner,
-    numBags: Int = -1,
-    useJackknife: Boolean = true,
-    biasLearner: Option[Learner] = None,
-    uncertaintyCalibration: Boolean = true,
-    disableBootstrap: Boolean = false
-) extends Learner {
-  require(
-    !(uncertaintyCalibration && disableBootstrap),
-    "Options uncertaintyCalibration and disableBootstrap are incompatible. At most one may be set true."
-  )
+sealed trait Bagger[T] extends Learner[T] {
+
+  def numBags: Int
+
+  def useJackknife: Boolean
+
+  def disableBootstrap: Boolean
+
+  def baseLearner: Learner[T]
 
   /**
-    * Draw with replacement from the training data for each model
+    * Draw with replacement from the training data for each model.
     *
     * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
     * @param rng          random number generator for reproducibility
-    * @return a model
+    * @return a training result containing the bagged model
     */
-  override def train(
-      trainingData: Seq[(Vector[Any], Any)],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): BaggedTrainingResult[Any] = {
-    /* Make sure the training data is the same size */
-    assert(trainingData.forall(trainingData.head._1.size == _._1.size))
+  override def train(trainingData: Seq[TrainingRow[T]], rng: Random = Random()): BaggedTrainingResult[T]
+
+  /** Bootstrap the training data to train an ensemble of models from the base learner. */
+  protected def trainEnsemble(trainingData: Seq[TrainingRow[T]], rng: Random): BaggedEnsemble[T] = {
+    // Make sure the training data is the same size
+    assert(trainingData.forall(trainingData.head.inputs.length == _.inputs.length))
     require(
-      trainingData.size >= Bagger.minimumTrainingSize,
-      s"We need to have at least ${Bagger.minimumTrainingSize} rows, only ${trainingData.size} given"
+      trainingData.length >= Bagger.minimumTrainingSize,
+      s"We need to have at least ${Bagger.minimumTrainingSize} rows, only ${trainingData.length} given."
     )
 
-    /* Use unit weights if none are specified */
-    val weightsActual = weights.getOrElse(Seq.fill(trainingData.size)(1.0))
-
-    /* Set default number of bags */
-    val actualBags = if (numBags > 0) {
-      numBags
-    } else {
-      trainingData.size
-    }
+    // Set default number of bags
+    val actualBags = if (numBags > 0) numBags else trainingData.length
 
     // We need enough bags such that the probability that the poisson draw is "valid" is at least 50%
     // Valid here means that for each training point, there is at least one tree that doesn't include it
     // The probability that the weights are valid is:
     // (1 - [(1 - 1/e)^{number of trees}])^{number of training points}
-    val minBags = Math.log(1 - Math.pow(2, -1.0 / trainingData.size)) / Math.log((Math.E - 1) / Math.E)
+    val minBags = math.log(1 - math.pow(2, -1.0 / trainingData.length)) / math.log((Math.E - 1) / math.E)
     require(
       !useJackknife || actualBags >= minBags,
-      s"Jackknife requires ${minBags} bags for ${trainingData.size} training rows, but only ${actualBags} given."
+      s"Jackknife requires $minBags bags for ${trainingData.length} training rows, but only $actualBags given."
     )
 
-    /* Compute the number of instances of each training row in each training sample */
-    val randBasis = breezeRandBasis(rng)
+    // Compute the number of instances of each training row in each training sample
+    val Nib = drawNib(actualBags, trainingData.length, rng)
+
+    // Learn the actual models in parallel
+    val indices = Nib.indices.toVector
+    val (models, importances) = rng
+      .zip(indices)
+      .par
+      .map {
+        case (thisRng, i) =>
+          val weightedTrainingData = Nib(i).zip(trainingData).map {
+            case (count, row) => row.mapWeight(_ * count.toDouble)
+          }
+          val meta = baseLearner.train(weightedTrainingData, thisRng)
+          (meta.model, meta.featureImportance)
+      }
+      .unzip
+
+    // Average the feature importance
+    val averageImportance = importances.reduce(Bagger.combineImportance).map(_.map(_ / importances.size))
+
+    BaggedEnsemble(models, Nib, averageImportance)
+  }
+
+  /**
+    * Compute the number of instances of each training row in each bag.
+    *
+    * @return the (# bags) x (# training rows) array of sample counts
+    */
+  private def drawNib(actualBags: Int, trainingSize: Int, rng: Random = Random()): Vector[Vector[Int]] = {
+    val randBasis = StatsUtils.breezeRandBasis(rng)
     val dist = new Poisson(1.0)(randBasis)
-    val Nib: Vector[Vector[Int]] = if (disableBootstrap) {
-      Vector.fill[Vector[Int]](actualBags)(Vector.fill[Int](trainingData.size)(1))
+    if (disableBootstrap) {
+      Vector.fill[Vector[Int]](actualBags)(Vector.fill[Int](trainingSize)(1))
     } else {
       Iterator
         .continually {
@@ -86,7 +91,7 @@ case class Bagger(
           // of non-zero training weights
           Iterator
             .continually {
-              Vector.fill(trainingData.size)(dist.draw())
+              Vector.fill(trainingSize)(dist.draw())
             }
             .filter(_.count(_ > 0) >= Bagger.minimumNonzeroWeightSize)
             .take(actualBags)
@@ -100,204 +105,100 @@ case class Bagger(
         }
         .next()
     }
-
-    // Learn the actual models in parallel
-    val indices = Nib.indices.toVector
-    val (models: ParSeq[Model[PredictionResult[Any]]], importances: ParSeq[Option[Vector[Double]]]) =
-      rng
-        .zip(indices)
-        .par
-        .map {
-          case (thisRng, i) =>
-            val sampleWeights = Nib(i).zip(weightsActual).map(p => p._1.toDouble * p._2)
-            val meta = method.train(trainingData.toVector, Some(sampleWeights), thisRng)
-            (meta.getModel(), meta.getFeatureImportance())
-        }
-        .unzip
-
-    // Average the feature importances
-    val averageImportance: Option[Vector[Double]] = importances
-      .reduce(Bagger.combineImportance)
-      .map(_.map(_ / importances.size))
-
-    /* Wrap the models in a BaggedModel object */
-    val helper = BaggerHelper(models, trainingData, Nib, useJackknife, uncertaintyCalibration)
-    val biasModel = if (biasLearner.isDefined && helper.oobErrors.nonEmpty && helper.isRegression) {
-      Some(
-        biasLearner.get.train(helper.biasTraining, rng = rng).getModel().asInstanceOf[Model[PredictionResult[Double]]]
-      )
-    } else None
-
-    if (helper.isRegression) {
-      new BaggedTrainingResult[Double](
-        models = models.asInstanceOf[ParSeq[Model[PredictionResult[Double]]]],
-        featureImportance = averageImportance,
-        Nib = Nib,
-        trainingData = trainingData,
-        useJackknife = useJackknife,
-        biasModel = biasModel,
-        rescale = helper.rescaleRatio,
-        disableBootstrap = disableBootstrap
-      )
-    } else {
-      new BaggedTrainingResult[Any](
-        models = models,
-        featureImportance = averageImportance,
-        Nib = Nib,
-        trainingData = trainingData,
-        useJackknife = useJackknife,
-        biasModel = biasModel,
-        rescale = helper.rescaleRatio,
-        disableBootstrap = disableBootstrap
-      )
-    }
   }
 }
 
-class BaggedTrainingResult[+T: ClassTag](
-    models: ParSeq[Model[PredictionResult[T]]],
-    featureImportance: Option[Vector[Double]],
-    Nib: Vector[Vector[Int]],
-    trainingData: Seq[(Vector[Any], Any)],
-    useJackknife: Boolean,
-    biasModel: Option[Model[PredictionResult[T]]] = None,
-    rescale: Double = 1.0,
-    disableBootstrap: Boolean = false
-) extends TrainingResult {
-
-  lazy val NibT = Nib.transpose
-  lazy val model = new BaggedModel[T](models, Nib, useJackknife, biasModel, rescale, disableBootstrap)
-  lazy val rep = trainingData.find(_._2 != null).get._2
-  lazy val predictedVsActual = trainingData.zip(NibT).flatMap {
-    case ((f, l), nb) =>
-      val oob = if (disableBootstrap) {
-        models.zip(nb)
-      } else {
-        models.zip(nb).filter(_._2 == 0)
-      }
-
-      if (oob.isEmpty || l == null || (l.isInstanceOf[Double] && l.asInstanceOf[Double].isNaN)) {
-        Seq()
-      } else {
-        val predicted = l match {
-          case _: Double => oob.map(_._1.transform(Seq(f)).getExpected().head.asInstanceOf[Double]).sum / oob.size
-          case _: Any    => oob.map(_._1.transform(Seq(f)).getExpected().head).groupBy(identity).maxBy(_._2.size)._1
-        }
-        Seq((f, predicted, l))
-      }
-  }
-
-  lazy val loss: Double = rep match {
-    case _: Double => RegressionMetrics.RMSE(predictedVsActual.asInstanceOf[Seq[(Vector[Any], Double, Double)]])
-    case _: Any    => ClassificationMetrics.loss(predictedVsActual)
-  }
-
-  /**
-    * Average the influences across the ensemble of models
-    *
-    * @return feature influences as an array of doubles
-    */
-  override def getFeatureImportance(): Option[Vector[Double]] = featureImportance
-
-  override def getModel(): BaggedModel[Any] = model
+/**
+  * A bagger creates an ensemble of models by training the learner on random samples of the training data.
+  *
+  * @param baseLearner learner to train each model in the ensemble
+  * @param numBags number of base models to aggregate (default of -1 sets the number of models to the number of training rows)
+  * @param useJackknife whether to enable jackknife uncertainty estimate
+  * @param uncertaintyCalibration whether to enable empirical uncertainty calibration
+  * @param disableBootstrap whether to disable bootstrap (useful when `method` implements its own randomization)
+  * @param biasLearner learner to use for estimating bias
+  */
+case class RegressionBagger(
+    baseLearner: Learner[Double],
+    numBags: Int = -1,
+    useJackknife: Boolean = true,
+    uncertaintyCalibration: Boolean = true,
+    disableBootstrap: Boolean = false,
+    biasLearner: Option[Learner[Double]] = None
+) extends Bagger[Double] {
+  require(
+    !(uncertaintyCalibration && disableBootstrap),
+    "Options uncertaintyCalibration and disableBootstrap are incompatible. At most one may be set true."
+  )
 
-  override def getPredictedVsActual(): Option[Seq[(Vector[Any], Any, Any)]] = Some(predictedVsActual)
+  override def train(trainingData: Seq[TrainingRow[Double]], rng: Random): RegressionBaggerTrainingResult = {
+    // Train the ensemble of models from the data
+    val ensemble = trainEnsemble(trainingData, rng)
 
-  override def getLoss(): Option[Double] = {
-    if (predictedVsActual.nonEmpty) {
-      Some(loss)
-    } else {
-      None
+    // Compute uncertainty rescales and train the bias model (if present)
+    val helper = BaggerHelper(ensemble.models, trainingData, ensemble.Nib, useJackknife, uncertaintyCalibration)
+    val biasModel = biasLearner.collect {
+      case learner if helper.oobErrors.nonEmpty =>
+        learner.train(helper.biasTraining, rng = rng).model
     }
+
+    RegressionBaggerTrainingResult(
+      ensembleModels = ensemble.models,
+      Nib = ensemble.Nib,
+      trainingData = trainingData,
+      featureImportance = ensemble.averageImportance,
+      biasModel = biasModel,
+      rescaleRatio = helper.rescaleRatio,
+      disableBootstrap = disableBootstrap
+    )
   }
 }
 
 /**
-  * Container holding a parallel sequence of models and the sample counts used to train them
+  * A bagger creates an ensemble of models by training the learner on random samples of the training data
   *
-  * @param models in this bagged model
-  * @param Nib    training sample counts
+  * @param baseLearner learner to train each model in the ensemble
+  * @param numBags number of base models to aggregate (default of -1 sets the number of models to the number of training rows)
+  * @param useJackknife whether to enable jackknife uncertainty estimate
+  * @param disableBootstrap whether to disable bootstrap (useful when `method` implements its own randomization)
+  * @tparam T the type of label data
   */
-class BaggedModel[+T: ClassTag](
-    models: ParSeq[Model[PredictionResult[T]]],
-    Nib: Vector[Vector[Int]],
-    useJackknife: Boolean,
-    biasModel: Option[Model[PredictionResult[T]]] = None,
-    rescale: Double = 1.0,
+case class ClassificationBagger[T](
+    baseLearner: Learner[T],
+    numBags: Int = -1,
+    useJackknife: Boolean = true,
     disableBootstrap: Boolean = false
-) extends Model[BaggedResult[T]] {
-
-  /**
-    * Apply each model to the outputs and wrap them up
-    *
-    * @param inputs to apply the model to
-    * @return a predictionresult that includes uncertainties and scores
-    */
-  override def transform(inputs: Seq[Vector[Any]]): BaggedResult[T] = {
-    assert(inputs.forall(_.size == inputs.head.size))
-    val isRegression = classTag[T].runtimeClass == classOf[Double]
-
-    val bias = if (biasModel.isDefined) {
-      Some(biasModel.get.transform(inputs).getExpected().asInstanceOf[Seq[Double]])
-    } else {
-      None
-    }
-    val ensemblePredictions = models.map(model => model.transform(inputs)).seq
-
-    val res = if (inputs.size == 1 && isRegression) {
-      // In the special case of a single prediction on a real value, emit an optimized BaggedSingleResult
-      SinglePredictionBaggedResult(
-        ensemblePredictions.map(_.asInstanceOf[PredictionResult[Double]]),
-        Nib,
-        bias.map(_.head),
-        rescale,
-        disableBootstrap
-      )
-    } else if (isRegression) {
-      MultiPredictionBaggedResult(
-        ensemblePredictions.map(_.asInstanceOf[PredictionResult[Double]]),
-        Nib,
-        bias,
-        rescale,
-        disableBootstrap
-      )
-    } else {
-      BaggedClassificationResult(ensemblePredictions)
-    }
-
-    res.asInstanceOf[BaggedResult[T]]
+) extends Bagger[T] {
+
+  override def train(trainingData: Seq[TrainingRow[T]], rng: Random): ClassificationBaggerTrainingResult[T] = {
+    // Train the ensemble of models from the data
+    val ensemble = trainEnsemble(trainingData, rng)
+
+    ClassificationBaggerTrainingResult(
+      ensembleModels = ensemble.models,
+      Nib = ensemble.Nib,
+      trainingData = trainingData,
+      featureImportance = ensemble.averageImportance,
+      disableBootstrap = disableBootstrap
+    )
   }
+}
+
+object Bagger {
 
   /**
-    * Compute Shapley feature attributions for a given input
+    * Data class storing intermediate results from training a [[Bagger]]
     *
-    * @param input for which to compute feature attributions.
-    * @param omitFeatures feature indices to omit in computing Shapley values
-    * @return matrix of attributions for each feature and output
-    *         One row per feature, each of length equal to the output dimension.
-    *         The output dimension is 1 for single-task regression, or equal to the number of classification categories.
+    * @param models trained from each bag in the ensemble
+    * @param Nib array of sample counts of shape (# bags) x (# training rows)
+    * @param averageImportance across models in the ensemble
+    * @tparam T type of label data for the models
     */
-  override def shapley(input: Vector[Any], omitFeatures: Set[Int] = Set()): Option[DenseMatrix[Double]] = {
-    val ensembleShapley = models.map(model => model.shapley(input, omitFeatures))
-    if (ensembleShapley.head.isEmpty) {
-      None
-    } else {
-      assert(ensembleShapley.forall(x => x.isDefined))
-
-      def sumReducer(a: Option[DenseMatrix[Double]], b: Option[DenseMatrix[Double]]): Option[DenseMatrix[Double]] = {
-        Some(a.get + b.get)
-      }
-      val scale = 1.0 / ensembleShapley.length
-
-      Some(scale * ensembleShapley.reduce(sumReducer).get)
-    }
-  }
-
-  // Accessor useful for testing.
-  private[bags] def getModels(): ParSeq[Model[PredictionResult[T]]] = models
-}
-
-object Bagger {
+  protected[bags] case class BaggedEnsemble[+T](
+      models: ParVector[Model[T]],
+      Nib: Vector[Vector[Int]],
+      averageImportance: Option[Vector[Double]]
+  )
 
   /**
     * The minimum number of training rows in order to train a Bagger
@@ -314,11 +215,10 @@ object Bagger {
     */
   val minimumNonzeroWeightSize: Int = 4
 
-  private 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
+  /** Combine two optional feature importance vectors. */
+  def combineImportance(vec1: Option[Vector[Double]], vec2: Option[Vector[Double]]): Option[Vector[Double]] = {
+    vec1.zip(vec2).map {
+      case (v1, v2) => v1.zip(v2).map(p => p._1 + p._2)
     }
   }
 }
diff --git a/src/main/scala/io/citrine/lolo/bags/BaggerHelper.scala b/src/main/scala/io/citrine/lolo/bags/BaggerHelper.scala
index 3f6d92d6..9659fc8f 100644
--- a/src/main/scala/io/citrine/lolo/bags/BaggerHelper.scala
+++ b/src/main/scala/io/citrine/lolo/bags/BaggerHelper.scala
@@ -1,12 +1,12 @@
 package io.citrine.lolo.bags
 
 import io.citrine.lolo.stats.{MathUtils, StatsUtils}
-import io.citrine.lolo.{Model, PredictionResult, RegressionResult}
+import io.citrine.lolo.api.{Model, TrainingRow}
 
 import scala.collection.parallel.immutable.ParSeq
 
 /**
-  * Helper class to subsume shared functionality of Bagger and MultiTaskBagger.
+  * Helper class to subsume shared functionality of [[RegressionBagger]] and [[MultiTaskBagger]].
   *
   * @param models collection of trained models
   * @param trainingData on which models were trained
@@ -14,51 +14,35 @@ import scala.collection.parallel.immutable.ParSeq
   * @param useJackknife whether to use jackknife for uncertainty quantification
   * @param uncertaintyCalibration whether to apply empirical uncertainty calibration
   */
-protected case class BaggerHelper(
-    models: ParSeq[Model[PredictionResult[Any]]],
-    trainingData: Seq[(Vector[Any], Any)],
+protected[bags] case class BaggerHelper(
+    models: ParSeq[Model[Double]],
+    trainingData: Seq[TrainingRow[Double]],
     Nib: Vector[Vector[Int]],
     useJackknife: Boolean,
     uncertaintyCalibration: Boolean
 ) {
-  val isRegression: Boolean = trainingData.map { _._2 }.find { _ != null } match {
-    case Some(_: Double) => true
-    case Some(_: Any)    => false
-    case None            => throw new IllegalArgumentException(s"Unable to find a non-null label")
-  }
 
   /**
     * Seq of tuples containing training data paired with their out-of-bag residuals and uncertainty estimates.
     * Defined as lazy so it's only computed if needed for the ratio or bias learner calculations.
     */
   lazy val oobErrors: Seq[(Vector[Any], Double, Double)] = trainingData.indices.flatMap { idx =>
-    val oobModels =
-      models.zip(Nib.map(_(idx))).filter(_._2 == 0).map(_._1).asInstanceOf[ParSeq[Model[PredictionResult[Double]]]]
-    val label = trainingData(idx)._2
-    if (oobModels.size < 2 || label == null || (label.isInstanceOf[Double] && label.asInstanceOf[Double].isNaN)) {
+    val oobModels = models.zip(Nib.map(_(idx))).filter(_._2 == 0).map(_._1)
+    val label = trainingData(idx).label
+    if (oobModels.size < 2 || label.isNaN) {
       None
     } else {
-      val model = new BaggedModel(
-        oobModels,
-        Nib.filter {
-          _(idx) == 0
-        },
-        useJackknife
-      )
-      val predicted = model.transform(Seq(trainingData(idx)._1))
-      val error = predicted.getExpected().head - trainingData(idx)._2.asInstanceOf[Double]
-      val uncertainty = predicted match {
-        case x: RegressionResult => x.getStdDevObs().get.head
-        case _: Any =>
-          throw new UnsupportedOperationException("Computing oobErrors for classification is not supported.")
-      }
-      Some(trainingData(idx)._1, error, uncertainty)
+      val model = BaggedRegressionModel(oobModels, Nib.filter { _(idx) == 0 })
+      val predicted = model.transform(Seq(trainingData(idx).inputs))
+      val error = predicted.expected.head - trainingData(idx).label
+      val uncertainty = predicted.stdDevObs.get.head
+      Some(trainingData(idx).inputs, error, uncertainty)
     }
   }
 
   /** Uncertainty calibration ratio based on the OOB errors. */
-  val rescaleRatio: Double = if (uncertaintyCalibration && isRegression && useJackknife) {
-    val trainingLabels = trainingData.collect { case (_, x: Double) if !(x.isInfinite || x.isNaN) => x }
+  val rescaleRatio: Double = if (uncertaintyCalibration && useJackknife) {
+    val trainingLabels = trainingData.collect { case TrainingRow(_, y, _) if !(y.isInfinite || y.isNaN) => y }
     val zeroTolerance = StatsUtils.range(trainingLabels) / 1e12
     BaggerHelper.calculateRescaleRatio(oobErrors.map { case (_, e, u) => (e, u) }, zeroTolerance = zeroTolerance)
   } else {
@@ -66,19 +50,17 @@ protected case class BaggerHelper(
   }
   assert(!rescaleRatio.isNaN && !rescaleRatio.isInfinity, s"Uncertainty calibration ratio is not real: $rescaleRatio")
 
-  /**
-    * Data on which to train a bias learner.
-    */
-  lazy val biasTraining: Seq[(Vector[Any], Double)] = oobErrors.map {
+  /** Data on which to train a bias learner. */
+  lazy val biasTraining: Seq[TrainingRow[Double]] = oobErrors.map {
     case (f, e, u) =>
       // 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(e) - u * rescaleRatio, 0)
-      (f, bias)
+      TrainingRow(f, bias)
   }
 }
 
-object BaggerHelper {
+protected[bags] object BaggerHelper {
 
   /**
     * Calculate the uncertainty calibration ratio, which is the 68th percentile of error/uncertainty for the training points.
diff --git a/src/main/scala/io/citrine/lolo/bags/MultiTaskBagger.scala b/src/main/scala/io/citrine/lolo/bags/MultiTaskBagger.scala
index 5f8ada26..edf8b697 100644
--- a/src/main/scala/io/citrine/lolo/bags/MultiTaskBagger.scala
+++ b/src/main/scala/io/citrine/lolo/bags/MultiTaskBagger.scala
@@ -1,48 +1,56 @@
 package io.citrine.lolo.bags
 
 import breeze.stats.distributions.Poisson
-import io.citrine.lolo.stats.StatsUtils.breezeRandBasis
-import io.citrine.lolo.{Learner, Model, MultiTaskLearner, MultiTaskModel, MultiTaskTrainingResult, PredictionResult}
+import io.citrine.lolo.api.{
+  Learner,
+  Model,
+  MultiTaskLearner,
+  MultiTaskModel,
+  MultiTaskModelPredictionResult,
+  MultiTaskTrainingResult,
+  ParallelModelsPredictionResult,
+  PredictionResult,
+  TrainingRow
+}
+import io.citrine.lolo.stats.StatsUtils
 import io.citrine.random.Random
 import io.citrine.lolo.stats.metrics.{ClassificationMetrics, RegressionMetrics}
 
 import scala.collection.parallel.immutable.ParSeq
-import scala.collection.parallel.CollectionConverters.IterableIsParallelizable
+import scala.collection.parallel.CollectionConverters._
 
 /**
-  * Create an ensemble of multi-task models
+  * 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 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 biasLearner            learner to use for estimating bias
   */
 case class MultiTaskBagger(
     method: MultiTaskLearner,
     numBags: Int = -1,
     useJackknife: Boolean = true,
-    biasLearner: Option[Learner] = None,
-    uncertaintyCalibration: Boolean = true
+    uncertaintyCalibration: Boolean = true,
+    biasLearner: Option[Learner[Double]] = None
 ) extends MultiTaskLearner {
 
-  override def train(
-      trainingData: Seq[(Vector[Any], Vector[Any])],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): MultiTaskBaggedTrainingResult = {
-    val (inputs, labels) = trainingData.unzip
-    val numInputs = inputs.head.length
-    val numOutputs = labels.head.length
+  override def train(trainingData: Seq[TrainingRow[Vector[Any]]], rng: Random): MultiTaskBaggedTrainingResult = {
+    val numInputs = trainingData.head.inputs.length
+    val numOutputs = trainingData.head.label.length
     /* Make sure the training data are the same size */
-    assert(inputs.forall(numInputs == _.size))
-    assert(labels.forall(numOutputs == _.size))
+    assert(
+      trainingData.forall { row =>
+        row.inputs.length == numInputs && row.label.length == numOutputs
+      }
+    )
     assert(
       trainingData.size >= Bagger.minimumTrainingSize,
       s"We need to have at least ${Bagger.minimumTrainingSize} rows, only ${trainingData.size} given"
     )
     (0 until numOutputs).foreach { i =>
-      val numOutputValues = labels.count(row => validOutput(row(i)))
+      val numOutputValues = trainingData.count(row => validOutput(row.label(i)))
       assert(
         numOutputValues >= Bagger.minimumOutputCount,
         s"There must be at least ${Bagger.minimumOutputCount} data points for each output, but output $i only had $numOutputValues values."
@@ -53,76 +61,68 @@ case class MultiTaskBagger(
     val actualBags = if (numBags > 0) numBags else trainingData.size
 
     // Compute the number of instances of each training row in each training sample
-    val randBasis = breezeRandBasis(rng)
+    val randBasis = StatsUtils.breezeRandBasis(rng)
     val dist = new Poisson(1.0)(randBasis)
     val Nib: Vector[Vector[Int]] = Iterator
       .continually(Vector.fill(trainingData.size)(dist.draw()))
       .filter { suggestedCounts =>
-        val allOutputsRepresented = (0 until numOutputs).forall(i =>
-          labels.zip(suggestedCounts).exists { case (row, count) => validOutput(row(i)) && count > 0 }
-        )
+        val allOutputsRepresented = (0 until numOutputs).forall { i =>
+          trainingData.zip(suggestedCounts).exists { case (row, count) => validOutput(row.label(i)) && count > 0 }
+        }
         val minNonzeroWeights = suggestedCounts.count(_ > 0) >= Bagger.minimumNonzeroWeightSize
         allOutputsRepresented && minNonzeroWeights
       }
       .take(actualBags)
       .toVector
-    val weightsActual = weights.getOrElse(Seq.fill(trainingData.size)(1.0))
 
     val indices = Nib.indices.toVector
-    val (models: ParSeq[MultiTaskModel], importances: ParSeq[Option[Vector[Double]]]) =
+    val (ensembleModels, importances) =
       rng
         .zip(indices)
         .par
         .map {
           case (thisRng, i) =>
-            val sampleWeights = Nib(i).zip(weightsActual).map(p => p._1.toDouble * p._2)
-            val meta = method.train(trainingData, Some(sampleWeights), thisRng)
-            (meta.getModel(), meta.getFeatureImportance())
+            val weightedTrainingData = Nib(i).zip(trainingData).map {
+              case (count, row) => row.mapWeight(_ * count.toDouble)
+            }
+            val meta = method.train(weightedTrainingData, thisRng)
+            (meta.model, meta.featureImportance)
         }
         .unzip
 
     val averageImportance: Option[Vector[Double]] = importances
-      .reduce(combineImportance)
+      .reduce(Bagger.combineImportance)
       .map(_.map(_ / importances.size))
 
     // Get bias model and rescale ratio for each label
-    val (biasModels, ratios) = Seq
+    val (biasModels, rescaleRatios) = Seq
       .tabulate(numOutputs) { 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) {
-          Some(
-            biasLearner.get
-              .train(helper.biasTraining, rng = rng)
-              .getModel()
-              .asInstanceOf[Model[PredictionResult[Double]]]
-          )
-        } else None
-        (biasModel, helper.rescaleRatio)
+        val isRegression = ensembleModels.head.realLabels(i)
+        if (isRegression) {
+          val thisLabelModels = ensembleModels.map(_.models(i).asInstanceOf[Model[Double]])
+          val thisTrainingData = trainingData.map(_.mapLabel(vec => vec(i).asInstanceOf[Double]))
+          val helper = BaggerHelper(thisLabelModels, thisTrainingData, Nib, useJackknife, uncertaintyCalibration)
+          val biasModel = biasLearner.collect {
+            case learner if helper.oobErrors.nonEmpty =>
+              learner.train(helper.biasTraining, rng = rng).model
+          }
+          (biasModel, helper.rescaleRatio)
+        } else {
+          (None, 1.0) // Rescale not used for classification tasks, so just set a default of 1.0
+        }
       }
       .unzip
 
-    new MultiTaskBaggedTrainingResult(
-      models = models,
-      featureImportance = averageImportance,
+    MultiTaskBaggedTrainingResult(
+      ensembleModels = ensembleModels,
       Nib = Nib,
+      featureImportance = averageImportance,
       trainingData = trainingData,
-      useJackknife = useJackknife,
       biasModels = biasModels,
-      rescaleRatios = ratios
+      rescaleRatios = rescaleRatios
     )
   }
 
-  /** Combine two optional feature importance vectors. */
-  def combineImportance(v1: Option[Vector[Double]], v2: Option[Vector[Double]]): Option[Vector[Double]] = {
-    (v1, v2) match {
-      case (Some(v1: Vector[Double]), Some(v2: Vector[Double])) => Some(v1.zip(v2).map(p => p._1 + p._2))
-      case _                                                    => None
-    }
-  }
-
   /** Flag NaNs and nulls. */
   private def validOutput(x: Any): Boolean = {
     Option(x) match {
@@ -136,42 +136,40 @@ case class MultiTaskBagger(
 /**
   * 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
+  * @param ensembleModels    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 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]],
+case class MultiTaskBaggedTrainingResult(
+    ensembleModels: ParSeq[MultiTaskModel],
     Nib: Vector[Vector[Int]],
-    trainingData: Seq[(Vector[Any], Seq[Any])],
-    useJackknife: Boolean,
-    biasModels: Seq[Option[Model[PredictionResult[Double]]]],
+    trainingData: Seq[TrainingRow[Vector[Any]]],
+    override val featureImportance: Option[Vector[Double]],
+    biasModels: Seq[Option[Model[Double]]],
     rescaleRatios: Seq[Double]
 ) extends MultiTaskTrainingResult {
 
-  lazy val model = new MultiTaskBaggedModel(models, Nib, useJackknife, biasModels, rescaleRatios)
+  override lazy val model: MultiTaskBaggedModel = MultiTaskBaggedModel(ensembleModels, Nib, biasModels, rescaleRatios)
 
   // 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]])] =
+  override lazy val predictedVsActual: Option[Seq[(Vector[Any], Vector[Option[Any]], Vector[Option[Any]])]] = Some(
     trainingData.zip(Nib.transpose).flatMap {
-      case ((features, labels), nb) =>
+      case (TrainingRow(features, labels, _), nb) =>
         // Bagged models that were not trained on this input
-        val oob = models.zip(nb).filter(_._2 == 0).map(_._1)
+        val oob = ensembleModels.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 oobPredictions = oob.map(_.transform(Seq(features)).expected.head)
           val predicted = oobPredictions.toVector.transpose.zipWithIndex.map {
-            case (predictions, labelIndex) if models.head.getRealLabels(labelIndex) =>
+            case (predictions, labelIndex) if ensembleModels.head.realLabels(labelIndex) =>
               predictions.asInstanceOf[Seq[Double]].sum / predictions.size
             case (predictions, _) => predictions.groupBy(identity).maxBy(_._2.size)._1
           }
@@ -186,101 +184,150 @@ class MultiTaskBaggedTrainingResult(
           Seq((features, optionPredicted, optionLabels))
         }
     }
+  )
 
-  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.lazyZip(allActual).lazyZip(models.head.getRealLabels).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.lazyZip(labelPredicted).lazyZip(labelActual).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)
+  override lazy val loss: Option[Double] = predictedVsActual.collect {
+    case pva if pva.nonEmpty =>
+      val allInputs = pva.map(_._1)
+      val allPredicted: Seq[Seq[Option[Any]]] = pva.map(_._2).transpose
+      val allActual: Seq[Seq[Option[Any]]] = pva.map(_._3).transpose
+      allPredicted
+        .lazyZip(allActual)
+        .lazyZip(ensembleModels.head.realLabels)
+        .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 labelPVA = allInputs.lazyZip(labelPredicted).lazyZip(labelActual).flatMap {
+              case (input, Some(p), Some(a)) => Some((input, p, a))
+              case _                         => None
+            }
+            if (isReal) {
+              RegressionMetrics.RMSE(labelPVA.asInstanceOf[Seq[(Vector[Any], Double, Double)]])
+            } else {
+              ClassificationMetrics.loss(labelPVA)
+            }
         }
-    }
-  }.sum
-
-  override def getFeatureImportance(): Option[Vector[Double]] = featureImportance
-
-  override def getModel(): MultiTaskBaggedModel = model
+        .sum
+  }
 
-  override def getModels(): Seq[Model[PredictionResult[Any]]] = {
-    val realLabels: Seq[Boolean] = models.head.getRealLabels
+  override def models: Seq[Model[Any]] = {
+    val realLabels: Seq[Boolean] = ensembleModels.head.realLabels
     realLabels.zipWithIndex.map {
       case (isReal: Boolean, i: Int) =>
-        val thisLabelModels = models.map(_.getModels(i))
+        val thisLabelModels = ensembleModels.map(_.models(i))
         if (isReal) {
-          new BaggedModel[Double](
-            thisLabelModels.asInstanceOf[ParSeq[Model[PredictionResult[Double]]]],
-            Nib,
-            useJackknife,
-            biasModels(i),
-            rescaleRatios(i)
+          BaggedRegressionModel(
+            thisLabelModels.asInstanceOf[ParSeq[Model[Double]]],
+            Nib = Nib,
+            rescaleRatio = rescaleRatios(i),
+            biasModel = biasModels(i)
           )
         } else {
-          new BaggedModel[Any](
-            thisLabelModels.asInstanceOf[ParSeq[Model[PredictionResult[Any]]]],
-            Nib,
-            useJackknife,
-            biasModels(i),
-            rescaleRatios(i)
-          )
+          BaggedClassificationModel(thisLabelModels)
         }
     }
   }
-
-  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
+  * @param ensembleModels sequence of multi-models, one for each bag
+  * @param Nib            matrix representing number of times each training datum appears in each bag
+  * @param biasModels     sequence of optional bias-correction models, one for each label
   * @param rescaleRatios  sequence of uncertainty calibration ratios for each label
   */
-class MultiTaskBaggedModel(
-    models: ParSeq[MultiTaskModel],
+case class MultiTaskBaggedModel(
+    ensembleModels: ParSeq[MultiTaskModel],
     Nib: Vector[Vector[Int]],
-    useJackknife: Boolean,
-    biasModels: Seq[Option[Model[PredictionResult[Double]]]],
+    biasModels: Seq[Option[Model[Double]]],
     rescaleRatios: Seq[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),
-        rescaleRatios(i)
+  override val numLabels: Int = ensembleModels.head.numLabels
+
+  override lazy val models: Vector[BaggedModel[Any]] = Vector.tabulate(numLabels) { i =>
+    val thisLabelsModels = ensembleModels.map(_.models(i))
+    if (realLabels(i)) {
+      BaggedRegressionModel(
+        thisLabelsModels.asInstanceOf[ParSeq[Model[Double]]],
+        Nib = Nib,
+        rescaleRatio = rescaleRatios(i),
+        biasModel = biasModels(i)
       )
     } else {
-      new BaggedModel(thisLabelsModels, Nib, useJackknife, biasModels(i), rescaleRatios(i))
+      BaggedClassificationModel(thisLabelsModels)
     }
   }
 
-  override def transform(inputs: Seq[Vector[Any]]) =
-    MultiTaskBaggedResult(groupedModels.map(_.transform(inputs)), getRealLabels, Nib)
+  override def transform(inputs: Seq[Vector[Any]]): MultiTaskBaggedPrediction =
+    MultiTaskBaggedPrediction(models.map(_.transform(inputs)), realLabels)
+
+  override def realLabels: Seq[Boolean] = ensembleModels.head.realLabels
+}
+
+/**
+  * Container with model-wise predictions for each label and the machinery to compute (co)variance.
+  *
+  * @param labelPredictions bagged prediction results for each label
+  * @param realLabels       a boolean sequence indicating which labels are real-valued
+  */
+case class MultiTaskBaggedPrediction(labelPredictions: Vector[BaggedPrediction[Any]], realLabels: Seq[Boolean])
+    extends BaggedPrediction[Vector[Any]]
+    with MultiTaskModelPredictionResult {
 
-  override val numLabels: Int = models.head.numLabels
+  override def numPredictions: Int = labelPredictions.head.numPredictions
 
-  override def getRealLabels: Seq[Boolean] = models.head.getRealLabels
+  override def expected: Seq[Vector[Any]] = labelPredictions.map(_.expected).transpose
+
+  override def ensemblePredictions: Seq[PredictionResult[Vector[Any]]] =
+    labelPredictions
+      .map(_.ensemblePredictions.map(_.expected))
+      .transpose
+      .map(x => 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 uncertainty(observational: Boolean = true): Option[Seq[Seq[Any]]] = {
+    Some(labelPredictions.map { predictionResult =>
+      predictionResult.uncertainty(observational) match {
+        case Some(value) => value
+        case None        => Seq.fill(numPredictions)(Double.NaN)
+      }
+    }.transpose)
+  }
+
+  override def uncertaintyCorrelation(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 =
+      labelPredictions(i).ensemblePredictions.map(_.expected).transpose.asInstanceOf[Seq[Seq[Double]]]
+    val baggedPredictionsJ =
+      labelPredictions(j).ensemblePredictions.map(_.expected).transpose.asInstanceOf[Seq[Seq[Double]]]
+    baggedPredictionsI.zip(baggedPredictionsJ).map {
+      case (bagsI, bagsJ) =>
+        StatsUtils.correlation(bagsI, bagsJ)
+    }
+  }
 
-  override def getModels: Seq[BaggedModel[Any]] = groupedModels
+  /**
+    * 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)
 }
diff --git a/src/main/scala/io/citrine/lolo/hypers/GridHyperOptimizer.scala b/src/main/scala/io/citrine/lolo/hypers/GridHyperOptimizer.scala
index 68fac3dc..7de69216 100644
--- a/src/main/scala/io/citrine/lolo/hypers/GridHyperOptimizer.scala
+++ b/src/main/scala/io/citrine/lolo/hypers/GridHyperOptimizer.scala
@@ -1,13 +1,9 @@
 package io.citrine.lolo.hypers
 
-import io.citrine.lolo.Learner
+import io.citrine.lolo.api.{Learner, TrainingRow}
 import io.citrine.random.Random
 
-/**
-  * Brute force search over the grid of hypers
-  *
-  * Created by maxhutch on 12/7/16.
-  */
+/** Brute force search over the grid of hypers. */
 case class GridHyperOptimizer() extends HyperOptimizer {
 
   /**
@@ -17,10 +13,10 @@ case class GridHyperOptimizer() extends HyperOptimizer {
     * @param numIterations ignored, since this is a brute force search
     * @return the best hyper map found in the search space
     */
-  override def optimize(
-      trainingData: Seq[(Vector[Any], Any)],
+  override def optimize[T](
+      trainingData: Seq[TrainingRow[T]],
       numIterations: Int = 1,
-      builder: Map[String, Any] => Learner,
+      builder: Map[String, Any] => Learner[T],
       rng: Random
   ): (Map[String, Any], Double) = {
     var best: Map[String, Any] = Map()
@@ -43,16 +39,16 @@ case class GridHyperOptimizer() extends HyperOptimizer {
       // Set up a learner with these parameters and compute the loss
       val testLearner = builder(testHypers)
       val res = testLearner.train(trainingData, rng = rng)
-      if (res.getLoss().isEmpty) {
-        throw new IllegalArgumentException("Trying to optimize hyper-paramters for a learner without getLoss")
+      if (res.loss.isEmpty) {
+        throw new IllegalArgumentException("Trying to optimize hyper-parameters for a learner without getLoss")
       }
-      val thisLoss = res.getLoss().get
+      val thisLoss = res.loss.get
 
       /* Save if it is an improvement */
       if (thisLoss < loss) {
         best = testHypers
         loss = thisLoss
-        println(s"Improved the loss to ${loss} with ${best}")
+        println(s"Improved the loss to $loss with $best")
       }
     }
     (best, loss)
diff --git a/src/main/scala/io/citrine/lolo/hypers/HyperOptimizer.scala b/src/main/scala/io/citrine/lolo/hypers/HyperOptimizer.scala
index 05c61d2a..820d22c1 100644
--- a/src/main/scala/io/citrine/lolo/hypers/HyperOptimizer.scala
+++ b/src/main/scala/io/citrine/lolo/hypers/HyperOptimizer.scala
@@ -1,13 +1,12 @@
 package io.citrine.lolo.hypers
 
-import io.citrine.lolo.Learner
+import io.citrine.lolo.api.{Learner, TrainingRow}
 import io.citrine.random.Random
 
 /**
   * Base class for hyperparameter optimizers
   *
   * They take a range of hypers as a Map[(String, Seq[Any])] and output the best map and loss
-  * Created by maxhutch on 12/8/16.
   */
 abstract class HyperOptimizer() {
 
@@ -33,10 +32,10 @@ abstract class HyperOptimizer() {
     * @param numIterations to take before terminating
     * @return the best hyper map found in give iterations and the corresponding loss
     */
-  def optimize(
-      trainingData: Seq[(Vector[Any], Any)],
+  def optimize[T](
+      trainingData: Seq[TrainingRow[T]],
       numIterations: Int = 8,
-      builder: Map[String, Any] => Learner,
+      builder: Map[String, Any] => Learner[T],
       rng: Random = Random()
   ): (Map[String, Any], Double)
 }
diff --git a/src/main/scala/io/citrine/lolo/hypers/RandomHyperOptimizer.scala b/src/main/scala/io/citrine/lolo/hypers/RandomHyperOptimizer.scala
index 35a2f202..57cab8d3 100644
--- a/src/main/scala/io/citrine/lolo/hypers/RandomHyperOptimizer.scala
+++ b/src/main/scala/io/citrine/lolo/hypers/RandomHyperOptimizer.scala
@@ -1,6 +1,6 @@
 package io.citrine.lolo.hypers
 
-import io.citrine.lolo.Learner
+import io.citrine.lolo.api.{Learner, TrainingRow}
 import io.citrine.random.Random
 
 /**
@@ -8,7 +8,6 @@ import io.citrine.random.Random
   *
   * This optimizer can be evaluated multiple times and will persist the best results across those
   * calls.
-  * Created by maxhutch on 12/7/16.
   */
 class RandomHyperOptimizer() extends HyperOptimizer {
 
@@ -16,7 +15,7 @@ class RandomHyperOptimizer() extends HyperOptimizer {
   var best: Map[String, Any] = Map()
 
   /** Likewise with the lowest loss */
-  var loss = Double.MaxValue
+  var loss: Double = Double.MaxValue
 
   /**
     * Search over numIterations random draws for the search space
@@ -25,29 +24,29 @@ class RandomHyperOptimizer() extends HyperOptimizer {
     * @param numIterations number of draws to take
     * @return the best hyper map found in give iterations and the corresponding loss
     */
-  override def optimize(
-      trainingData: Seq[(Vector[Any], Any)],
+  override def optimize[T](
+      trainingData: Seq[TrainingRow[T]],
       numIterations: Int,
-      builder: Map[String, Any] => Learner,
+      builder: Map[String, Any] => Learner[T],
       rng: Random
   ): (Map[String, Any], Double) = {
     /* Just draw numIteration times */
-    (0 until numIterations).foreach { i =>
+    (0 until numIterations).foreach { _ =>
       val testHypers = hyperGrids.map {
         case (n, v) =>
           n -> rng.shuffle(v).head
       }
       val testLearner = builder(testHypers)
       val res = testLearner.train(trainingData, rng = rng)
-      if (res.getLoss().isEmpty) {
+      if (res.loss.isEmpty) {
         throw new IllegalArgumentException("Trying to optimize hyper-paramters for a learner without getLoss")
       }
-      val thisLoss = res.getLoss().get
+      val thisLoss = res.loss.get
       /* Keep track of the best */
       if (thisLoss < loss) {
         best = testHypers
         loss = thisLoss
-        println(s"Improved the loss to ${loss} with ${best}")
+        println(s"Improved the loss to $loss with $best")
       }
     }
     (best, loss)
diff --git a/src/main/scala/io/citrine/lolo/learners/ExtraRandomTrees.scala b/src/main/scala/io/citrine/lolo/learners/ExtraRandomTrees.scala
index 095fba71..6bb529d9 100644
--- a/src/main/scala/io/citrine/lolo/learners/ExtraRandomTrees.scala
+++ b/src/main/scala/io/citrine/lolo/learners/ExtraRandomTrees.scala
@@ -1,15 +1,15 @@
 package io.citrine.lolo.learners
 
+import io.citrine.lolo.api.{Learner, TrainingResult, TrainingRow}
 import io.citrine.random.Random
-import io.citrine.lolo.bags.Bagger
-import io.citrine.lolo.transformers.FeatureRotator
+import io.citrine.lolo.bags.{ClassificationBagger, RegressionBagger}
+import io.citrine.lolo.transformers.rotator.FeatureRotator
 import io.citrine.lolo.trees.classification.ClassificationTreeLearner
 import io.citrine.lolo.trees.regression.RegressionTreeLearner
 import io.citrine.lolo.trees.splits.{ExtraRandomClassificationSplitter, ExtraRandomRegressionSplitter}
-import io.citrine.lolo.{Learner, TrainingResult}
 
 /**
-  * Extremely randomized tree ensemble
+  * Extremely randomized tree ensemble for regression.
   *
   * This is based on Geurts, P., Ernst, D. & Wehenkel, L. Extremely randomized trees. Mach Learn 63, 3–42 (2006).
   * https://doi.org/10.1007/s10994-006-6226-1.
@@ -26,71 +26,84 @@ import io.citrine.lolo.{Learner, TrainingResult}
   * @param disableBootstrap whether to disable bootstrap (default: true)
   * @param randomlyRotateFeatures whether to randomly rotate real features for each tree in the forest (default: false)
   */
-case class ExtraRandomTrees(
+case class ExtraRandomTreesRegressor(
     numTrees: Int = -1,
     useJackknife: Boolean = false,
-    biasLearner: Option[Learner] = None,
-    leafLearner: Option[Learner] = None,
+    biasLearner: Option[Learner[Double]] = None,
+    leafLearner: Option[Learner[Double]] = None,
     subsetStrategy: Any = "auto",
     minLeafInstances: Int = 1,
     maxDepth: Int = Integer.MAX_VALUE,
     uncertaintyCalibration: Boolean = false,
     disableBootstrap: Boolean = true,
     randomlyRotateFeatures: Boolean = false
-) extends Learner {
+) extends Learner[Double] {
 
-  /**
-    * Train an extremely randomized tree ensemble model
-    *
-    * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
-    * @param rng          random number generator for reproducibility
-    * @return training result containing a model
-    */
-  override def train(
-      trainingData: Seq[(Vector[Any], Any)],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): TrainingResult = {
-    val repOutput = trainingData.head._2
-    val isRegression = repOutput.isInstanceOf[Double]
-    val numFeatures: Int = RandomForest.getNumFeatures(subsetStrategy, trainingData.head._1.size, isRegression)
+  override def train(trainingData: Seq[TrainingRow[Double]], rng: Random): TrainingResult[Double] = {
+    val numFeatures = RandomForest.getNumFeatures(subsetStrategy, trainingData.head.inputs.length, isRegression = true)
 
-    repOutput match {
-      case _: Double =>
-        val DTLearner = RegressionTreeLearner(
-          leafLearner = leafLearner,
-          numFeatures = numFeatures,
-          minLeafInstances = minLeafInstances,
-          maxDepth = maxDepth,
-          splitter = ExtraRandomRegressionSplitter()
-        )
-        val bagger = Bagger(
-          if (randomlyRotateFeatures) FeatureRotator(DTLearner) else DTLearner,
-          numBags = numTrees,
-          useJackknife = useJackknife,
-          biasLearner = biasLearner,
-          uncertaintyCalibration = uncertaintyCalibration,
-          disableBootstrap = disableBootstrap
-        )
-        bagger.train(trainingData, weights, rng)
-      case _: Any =>
-        val DTLearner = ClassificationTreeLearner(
-          leafLearner = leafLearner,
-          numFeatures = numFeatures,
-          minLeafInstances = minLeafInstances,
-          maxDepth = maxDepth,
-          splitter = ExtraRandomClassificationSplitter()
-        )
-        val bagger = Bagger(
-          if (randomlyRotateFeatures) FeatureRotator(DTLearner) else DTLearner,
-          numBags = numTrees,
-          useJackknife = useJackknife,
-          biasLearner = biasLearner,
-          uncertaintyCalibration = uncertaintyCalibration,
-          disableBootstrap = disableBootstrap
-        )
-        bagger.train(trainingData, weights, rng)
-    }
+    val DTLearner = RegressionTreeLearner(
+      leafLearner = leafLearner,
+      numFeatures = numFeatures,
+      minLeafInstances = minLeafInstances,
+      maxDepth = maxDepth,
+      splitter = ExtraRandomRegressionSplitter()
+    )
+    val bagger = RegressionBagger(
+      if (randomlyRotateFeatures) FeatureRotator(DTLearner) else DTLearner,
+      numBags = numTrees,
+      useJackknife = useJackknife,
+      biasLearner = biasLearner,
+      uncertaintyCalibration = uncertaintyCalibration,
+      disableBootstrap = disableBootstrap
+    )
+    bagger.train(trainingData, rng)
+  }
+}
+
+/**
+  * Extremely randomized tree ensemble for classification.
+  *
+  * This is based on Geurts, P., Ernst, D. & Wehenkel, L. Extremely randomized trees. Mach Learn 63, 3–42 (2006).
+  * https://doi.org/10.1007/s10994-006-6226-1.
+  *
+  * @param numTrees         number of trees to use (-1 => number of training instances)
+  * @param useJackknife     whether to use jackknife based variance estimates
+  * @param leafLearner      learner to use at the leaves of the trees
+  * @param subsetStrategy   for random feature selection at each split
+  *                         (auto => all features for regression)
+  * @param minLeafInstances minimum number of instances per leave in each tree
+  * @param maxDepth         maximum depth of each tree in the forest (default: unlimited)
+  * @param disableBootstrap whether to disable bootstrap (default: true)
+  * @param randomlyRotateFeatures whether to randomly rotate real features for each tree in the forest (default: false)
+  */
+case class ExtraRandomTreesClassifier(
+    numTrees: Int = -1,
+    useJackknife: Boolean = false,
+    leafLearner: Option[Learner[Char]] = None,
+    subsetStrategy: Any = "auto",
+    minLeafInstances: Int = 1,
+    maxDepth: Int = Integer.MAX_VALUE,
+    disableBootstrap: Boolean = true,
+    randomlyRotateFeatures: Boolean = false
+) extends Learner[Any] {
+
+  override def train(trainingData: Seq[TrainingRow[Any]], rng: Random): TrainingResult[Any] = {
+    val numFeatures = RandomForest.getNumFeatures(subsetStrategy, trainingData.head.inputs.length, isRegression = false)
+
+    val DTLearner = ClassificationTreeLearner(
+      leafLearner = leafLearner,
+      numFeatures = numFeatures,
+      minLeafInstances = minLeafInstances,
+      maxDepth = maxDepth,
+      splitter = ExtraRandomClassificationSplitter()
+    )
+    val bagger = ClassificationBagger(
+      if (randomlyRotateFeatures) FeatureRotator(DTLearner) else DTLearner,
+      numBags = numTrees,
+      useJackknife = useJackknife,
+      disableBootstrap = disableBootstrap
+    )
+    bagger.train(trainingData, rng)
   }
 }
diff --git a/src/main/scala/io/citrine/lolo/learners/RandomForest.scala b/src/main/scala/io/citrine/lolo/learners/RandomForest.scala
index 14c20c8a..003ad394 100644
--- a/src/main/scala/io/citrine/lolo/learners/RandomForest.scala
+++ b/src/main/scala/io/citrine/lolo/learners/RandomForest.scala
@@ -1,17 +1,19 @@
 package io.citrine.lolo.learners
 
+import io.citrine.lolo.api.{Learner, MultiTaskLearner, MultiTaskTrainingResult, TrainingResult, TrainingRow}
 import io.citrine.random.Random
-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.bags.{ClassificationBagger, MultiTaskBagger, RegressionBagger}
+import io.citrine.lolo.transformers.rotator.{FeatureRotator, MultiTaskFeatureRotator}
+import io.citrine.lolo.transformers.standardizer.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, MultiTaskSplitter, RegressionSplitter}
-import io.citrine.lolo.{Learner, TrainingResult}
+
+import scala.annotation.tailrec
 
 /**
-  * Standard random forest as a wrapper around bagged decision trees
+  * Standard random forest regressor as a wrapper around bagged decision trees.
   *
   * @param numTrees       number of trees to use (-1 => number of training instances)
   * @param useJackknife   whether to use jackknife based variance estimates
@@ -25,101 +27,136 @@ import io.citrine.lolo.{Learner, TrainingResult}
   * @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)
   */
-case class RandomForest(
+case class RandomForestRegressor(
     numTrees: Int = -1,
     useJackknife: Boolean = true,
-    biasLearner: Option[Learner] = None,
-    leafLearner: Option[Learner] = None,
+    biasLearner: Option[Learner[Double]] = None,
+    leafLearner: Option[Learner[Double]] = None,
     subsetStrategy: Any = "auto",
     minLeafInstances: Int = 1,
     maxDepth: Int = Integer.MAX_VALUE,
     uncertaintyCalibration: Boolean = true,
     randomizePivotLocation: Boolean = false,
     randomlyRotateFeatures: Boolean = false
-) extends Learner {
+) extends Learner[Double] {
 
-  /**
-    * Train a random forest model
-    *
-    * If the training label is a Double then this is a regression forest.
-    * If the training label is another primtiive then this is a classification forest.
-    * If the training label is a sequence then this is a multitask forest.
-    * Options like the number of trees are set via the constructor.
-    *
-    * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
-    * @param rng          random number generator for reproducibility
-    * @return training result containing a model
-    */
-  override def train(
-      trainingData: Seq[(Vector[Any], Any)],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): TrainingResult = {
-    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)
+  override def train(trainingData: Seq[TrainingRow[Double]], rng: Random): TrainingResult[Double] = {
+    val numFeatures = RandomForest.getNumFeatures(subsetStrategy, trainingData.head.inputs.length, isRegression = true)
 
-    repOutput match {
-      case _: Double =>
-        val DTLearner = RegressionTreeLearner(
-          leafLearner = leafLearner,
-          numFeatures = numFeatures,
-          minLeafInstances = minLeafInstances,
-          maxDepth = maxDepth,
-          splitter = RegressionSplitter(randomizePivotLocation)
-        )
+    val DTLearner = RegressionTreeLearner(
+      leafLearner = leafLearner,
+      numFeatures = numFeatures,
+      minLeafInstances = minLeafInstances,
+      maxDepth = maxDepth,
+      splitter = RegressionSplitter(randomizePivotLocation)
+    )
+    val bagger = RegressionBagger(
+      if (randomlyRotateFeatures) FeatureRotator(DTLearner) else DTLearner,
+      numBags = numTrees,
+      useJackknife = useJackknife,
+      biasLearner = biasLearner,
+      uncertaintyCalibration = uncertaintyCalibration
+    )
+    bagger.train(trainingData, rng)
+  }
+}
 
-        val bagger = Bagger(
-          if (randomlyRotateFeatures) FeatureRotator(DTLearner) else DTLearner,
-          numBags = numTrees,
-          useJackknife = useJackknife,
-          biasLearner = biasLearner,
-          uncertaintyCalibration = uncertaintyCalibration
-        )
-        bagger.train(trainingData, weights, rng)
-      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,
-            splitter = MultiTaskSplitter(randomizePivotLocation)
-          )
-        )
-        val bagger = MultiTaskBagger(
-          if (randomlyRotateFeatures) MultiTaskFeatureRotator(DTLearner) else DTLearner,
-          numBags = numTrees,
-          useJackknife = useJackknife,
-          biasLearner = biasLearner,
-          uncertaintyCalibration = uncertaintyCalibration
-        )
-        bagger.train(trainingData.asInstanceOf[Seq[(Vector[Any], Vector[Any])]], weights, rng)
-      case _: Any =>
-        val DTLearner = ClassificationTreeLearner(
-          leafLearner = leafLearner,
-          numFeatures = numFeatures,
-          minLeafInstances = minLeafInstances,
-          maxDepth = maxDepth,
-          splitter = ClassificationSplitter(randomizePivotLocation)
-        )
-        val bagger = Bagger(
-          if (randomlyRotateFeatures) FeatureRotator(DTLearner) else DTLearner,
-          numBags = numTrees
-        )
-        bagger.train(trainingData, weights, rng)
-    }
+/**
+  * Standard random forest classifier as a wrapper around bagged decision trees.
+  *
+  * @param numTrees       number of trees to use (-1 => number of training instances)
+  * @param useJackknife   whether to use jackknife based variance estimates
+  * @param leafLearner    learner to use at the leaves of the trees
+  * @param subsetStrategy for random feature selection at each split
+  *                       (auto => all fetures for regression, sqrt for classification)
+  * @param minLeafInstances minimum number of instances per leave in each tree
+  * @param maxDepth       maximum depth of each tree in the forest (default: unlimited)
+  * @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)
+  */
+case class RandomForestClassifier(
+    numTrees: Int = -1,
+    useJackknife: Boolean = true,
+    leafLearner: Option[Learner[Char]] = None,
+    subsetStrategy: Any = "auto",
+    minLeafInstances: Int = 1,
+    maxDepth: Int = Integer.MAX_VALUE,
+    randomizePivotLocation: Boolean = false,
+    randomlyRotateFeatures: Boolean = false
+) extends Learner[Any] {
+
+  override def train(trainingData: Seq[TrainingRow[Any]], rng: Random): TrainingResult[Any] = {
+    val numFeatures = RandomForest.getNumFeatures(subsetStrategy, trainingData.head.inputs.length, isRegression = true)
+
+    val DTLearner = ClassificationTreeLearner(
+      leafLearner = leafLearner,
+      numFeatures = numFeatures,
+      minLeafInstances = minLeafInstances,
+      maxDepth = maxDepth,
+      splitter = ClassificationSplitter(randomizePivotLocation)
+    )
+    val bagger = ClassificationBagger(
+      if (randomlyRotateFeatures) FeatureRotator(DTLearner) else DTLearner,
+      numBags = numTrees,
+      useJackknife = useJackknife
+    )
+    bagger.train(trainingData, rng)
+  }
+}
+
+/**
+  * Multi-task, mixed-task random forest model as a wrapper around bagged decision trees.
+  *
+  * @param numTrees         number of trees to use (-1 => number of training instances)
+  * @param useJackknife     whether to use jackknife based variance estimates
+  * @param biasLearner      learner to model bias (absolute residual)
+  * @param subsetStrategy   for random feature selection at each split
+  *                         (auto => all fetures for regression, sqrt for classification)
+  * @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 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)
+  */
+case class MultiTaskRandomForest(
+    numTrees: Int = -1,
+    useJackknife: Boolean = true,
+    biasLearner: Option[Learner[Double]] = None,
+    subsetStrategy: Any = "auto",
+    minLeafInstances: Int = 1,
+    maxDepth: Int = Integer.MAX_VALUE,
+    uncertaintyCalibration: Boolean = true,
+    randomizePivotLocation: Boolean = false,
+    randomlyRotateFeatures: Boolean = false
+) extends MultiTaskLearner {
+
+  override def train(trainingData: Seq[TrainingRow[Vector[Any]]], rng: Random): MultiTaskTrainingResult = {
+    val rep = trainingData.head.label
+    val hasRegression = rep.exists(_.isInstanceOf[Double])
+    val numFeatures = RandomForest.getNumFeatures(subsetStrategy, trainingData.head.inputs.length, hasRegression)
+
+    val DTLearner = MultiTaskStandardizer(
+      MultiTaskTreeLearner(
+        numFeatures = numFeatures,
+        maxDepth = maxDepth,
+        minLeafInstances = minLeafInstances,
+        splitter = MultiTaskSplitter(randomizePivotLocation)
+      )
+    )
+    val bagger = MultiTaskBagger(
+      if (randomlyRotateFeatures) MultiTaskFeatureRotator(DTLearner) else DTLearner,
+      numBags = numTrees,
+      useJackknife = useJackknife,
+      biasLearner = biasLearner,
+      uncertaintyCalibration = uncertaintyCalibration
+    )
+    bagger.train(trainingData, rng)
   }
 }
 
 object RandomForest {
 
+  @tailrec
   def getNumFeatures(subsetStrategy: Any, numTrainingFeatures: Int, isRegression: Boolean): Int = {
     subsetStrategy match {
       case x: String =>
@@ -138,5 +175,4 @@ object RandomForest {
         (numTrainingFeatures * x).toInt
     }
   }
-
 }
diff --git a/src/main/scala/io/citrine/lolo/linear/GuessTheMean.scala b/src/main/scala/io/citrine/lolo/linear/GuessTheMean.scala
index 8093e758..cb9e4523 100644
--- a/src/main/scala/io/citrine/lolo/linear/GuessTheMean.scala
+++ b/src/main/scala/io/citrine/lolo/linear/GuessTheMean.scala
@@ -1,53 +1,36 @@
 package io.citrine.lolo.linear
 
+import io.citrine.lolo.api.{Learner, Model, PredictionResult, TrainingResult, TrainingRow}
 import io.citrine.random.Random
-import io.citrine.lolo.{Learner, Model, PredictionResult, TrainingResult}
-
-/**
-  * Created by maxhutch on 11/15/16.
-  */
-case class GuessTheMeanLearner() extends Learner {
-
-  /**
-    * Train a model
-    *
-    * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
-    * @param rng          random number generator for reproducibility
-    * @return training result containing a model
-    */
-  override def train(
-      trainingData: Seq[(Vector[Any], Any)],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): TrainingResult = {
-    val data = trainingData.map(_._2).zip(weights.getOrElse(Seq.fill(trainingData.size)(1.0)))
-    val mean = data.head._1 match {
-      case _: Double => data.asInstanceOf[Seq[(Double, Double)]].map(p => p._1 * p._2).sum / data.map(_._2).sum
-      case _: Any    => rng.shuffle(data.groupBy(_._1).view.mapValues(_.map(_._2).sum).toSeq).maxBy(_._2)._1
-    }
-
-    new GuessTheMeanTrainingResult(new GuessTheMeanModel(mean))
-  }
-}
 
-class GuessTheMeanTrainingResult[T](model: GuessTheMeanModel[T]) extends TrainingResult {
-  override def getModel(): Model[GuessTheMeanResult[T]] = model
+/** Learner that computes the (weighted) mean value over real-valued labels. */
+case class GuessTheMeanLearner() extends Learner[Double] {
+
+  override def train(trainingData: Seq[TrainingRow[Double]], rng: Random): GuessTheMeanTrainingResult[Double] = {
+    val totalWeight = trainingData.map(_.weight).sum
+    val mean = trainingData.map { case TrainingRow(_, label, weight) => label * weight }.sum / totalWeight
+    GuessTheMeanTrainingResult(GuessTheMeanModel(mean))
+  }
 }
 
-class GuessTheMeanModel[T](mean: T) extends Model[GuessTheMeanResult[T]] {
+/** Learner that computes the (weighted) mode value over labels of type `T`. */
+case class GuessTheModeLearner[T]() extends Learner[T] {
 
-  def transform(inputs: Seq[Vector[Any]]): GuessTheMeanResult[T] = {
-    new GuessTheMeanResult(Seq.fill(inputs.size)(mean))
+  override def train(trainingData: Seq[TrainingRow[T]], rng: Random): GuessTheMeanTrainingResult[T] = {
+    val mode = rng
+      .shuffle { trainingData.groupBy(_.label).view.mapValues(_.map(_.weight).sum).toSeq }
+      .maxBy(_._2)
+      ._1
+    GuessTheMeanTrainingResult(GuessTheMeanModel(mode))
   }
 }
 
-class GuessTheMeanResult[T](result: Seq[T]) extends PredictionResult[T] {
+case class GuessTheMeanTrainingResult[+T](model: GuessTheMeanModel[T]) extends TrainingResult[T]
 
-  /**
-    * Get the expected values for this prediction
-    *
-    * @return expected value of each prediction
-    */
-  override def getExpected(): Seq[T] = result
+case class GuessTheMeanModel[+T](value: T) extends Model[T] {
+  def transform(inputs: Seq[Vector[Any]]): GuessTheMeanPrediction[T] = {
+    GuessTheMeanPrediction(Seq.fill(inputs.size)(value))
+  }
 }
+
+case class GuessTheMeanPrediction[+T](expected: Seq[T]) extends PredictionResult[T]
diff --git a/src/main/scala/io/citrine/lolo/linear/LinearRegression.scala b/src/main/scala/io/citrine/lolo/linear/LinearRegression.scala
index 87cedb1f..e724d4c6 100644
--- a/src/main/scala/io/citrine/lolo/linear/LinearRegression.scala
+++ b/src/main/scala/io/citrine/lolo/linear/LinearRegression.scala
@@ -1,41 +1,35 @@
 package io.citrine.lolo.linear
 
 import breeze.linalg.{diag, sum, DenseMatrix, DenseVector}
-import io.citrine.lolo.{Learner, Model, PredictionResult, TrainingResult}
+import io.citrine.lolo.api.{Learner, Model, PredictionResult, TrainingResult, TrainingRow}
 import io.citrine.random.Random
 import org.slf4j.LoggerFactory
 
 import scala.util.{Failure, Success, Try}
 
 /**
-  * Linear and ridge regression learner
-  *
-  * Created by maxhutch on 12/6/16.
+  * Linear ridge regression learner.
   *
+  * @param regParam regularization to use for the ridge model
   * @param fitIntercept whether to fit an intercept or not
   */
 case class LinearRegressionLearner(
     regParam: Option[Double] = None,
     fitIntercept: Boolean = true
-) extends Learner {
+) extends Learner[Double] {
 
   /**
     * Train a linear model via direct inversion.
     *
     * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
     * @param rng          random number generator for reproducibility
     * @return a model
     */
-  override def train(
-      trainingData: Seq[(Vector[Any], Any)],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): LinearRegressionTrainingResult = {
+  override def train(trainingData: Seq[TrainingRow[Double]], rng: Random): LinearRegressionTrainingResult = {
     val lambda = regParam.getOrElse(0.0)
     val regularized = lambda > 0.0
 
-    val (features, labels) = trainingData.unzip
+    val (features, labels, weights) = trainingData.map(_.asTuple).unzip3
     val rep = features.head
 
     /* Get the indices of the continuous features */
@@ -61,12 +55,12 @@ case class LinearRegressionLearner(
       new DenseMatrix(numFeatures, numSamples, featureArray)
     }
     val X = Xt.t
-    val y = new DenseVector(labels.map(_.asInstanceOf[Double]).toArray)
+    val y = new DenseVector(labels.toArray)
 
     /* Rescale data by weight matrix */
-    val weightsMatrix = weights.map(w => diag(new DenseVector(w.toArray)))
-    val Xw = weightsMatrix.map(W => W * X).getOrElse(X)
-    val yw = weightsMatrix.map(W => W * y).getOrElse(y)
+    val W = diag(new DenseVector(weights.toArray))
+    val Xw = W * X
+    val yw = W * y
 
     val (coefficients, intercept) = Try {
       // For a regularized/overdetermined problem, the LHS operator `A` is (potentially) invertible
@@ -91,8 +85,7 @@ case class LinearRegressionLearner(
         }
       case Failure(e) =>
         logger.warn(s"Encountered an exception solving normal equations: ${e.getLocalizedMessage}")
-        val totalWeight = weights.map(_.sum).getOrElse(numSamples.toDouble)
-        val mean = sum(yw) / totalWeight // weighted mean of training labels
+        val mean = sum(yw) / weights.sum // weighted mean of training labels
         if (fitIntercept) {
           (DenseVector.zeros[Double](numFeatures - 1), mean)
         } else {
@@ -107,8 +100,8 @@ case class LinearRegressionLearner(
       None
     }
 
-    val model = new LinearRegressionModel(coefficients, intercept, indices = indicesToModel)
-    new LinearRegressionTrainingResult(model)
+    val model = LinearRegressionModel(coefficients, intercept, indices = indicesToModel)
+    LinearRegressionTrainingResult(model)
   }
 
   private val logger = LoggerFactory.getLogger(getClass)
@@ -119,67 +112,64 @@ case class LinearRegressionLearner(
   *
   * @param model contained
   */
-class LinearRegressionTrainingResult(model: LinearRegressionModel) extends TrainingResult {
-
-  override def getModel(): LinearRegressionModel = model
+case class LinearRegressionTrainingResult(model: LinearRegressionModel) extends TrainingResult[Double] {
 
   /**
     * Get a measure of the importance of the model features
     *
     * @return feature influences as an array of doubles
     */
-  override def getFeatureImportance(): Option[Vector[Double]] = {
-    val beta: Vector[Double] = model.getBeta().map(Math.abs)
-    val renorm: Double = 1.0 / beta.sum
-    Some(beta.map(_ * renorm))
+  override def featureImportance: Option[Vector[Double]] = {
+    val importance = model.beta.map(math.abs)
+    val renorm = 1.0 / importance.sum
+    Some(importance.map(_ * renorm))
   }
 }
 
 /**
   * Linear regression model as a coefficient vector and intercept
   *
-  * @param beta      coefficient vector
+  * @param denseBeta coefficient vector (only for active indices in the feature vector)
   * @param intercept intercept
   * @param indices   optional indices from which to extract real features
   */
-class LinearRegressionModel(
-    beta: DenseVector[Double],
+case class LinearRegressionModel(
+    denseBeta: DenseVector[Double],
     intercept: Double,
     indices: Option[(Vector[Int], Int)] = None
-) extends Model[LinearRegressionResult] {
+) extends Model[Double] {
 
   /**
     * Apply the model to a seq of inputs
     *
     * @param inputs to apply the model to
-    * @return a predictionresult which includes, at least, the expected outputs
+    * @return a prediction result which includes, at least, the expected outputs
     */
-  override def transform(inputs: Seq[Vector[Any]]): LinearRegressionResult = {
+  override def transform(inputs: Seq[Vector[Any]]): LinearRegressionPrediction = {
     val filteredInputs = indices
-      .map { case (ind, size) => inputs.map(inp => ind.map(inp(_))) }
+      .map { case (ind, _) => inputs.map(inp => ind.map(inp(_))) }
       .getOrElse(inputs)
       .flatten
       .asInstanceOf[Seq[Double]]
     val inputMatrix = new DenseMatrix(filteredInputs.size / inputs.size, inputs.size, filteredInputs.toArray)
-    val resultVector = beta.t * inputMatrix + intercept
+    val resultVector = denseBeta.t * inputMatrix + intercept
     val result = resultVector.t.toArray.toSeq
-    val grad = getBeta()
-    new LinearRegressionResult(result, grad)
+    LinearRegressionPrediction(result, beta)
   }
 
   /**
     * Get the beta from the linear model \beta^T X = y
     * @return beta as a vector of double
     */
-  def getBeta(): Vector[Double] = {
+  lazy val beta: Vector[Double] = {
     indices
       .map {
         case (inds, size) =>
           val empty = DenseVector.zeros[Double](size)
-          inds.zipWithIndex.foreach { case (j, i) => empty(j) = beta(i) }
+          inds.zipWithIndex.foreach { case (j, i) => empty(j) = denseBeta(i) }
           empty
       }
-      .getOrElse(beta)
+      .getOrElse(denseBeta)
       .toScalaVector
   }
 }
@@ -187,22 +177,15 @@ class LinearRegressionModel(
 /**
   * Simple container around the result and coefficient array
   *
-  * @param values computed from the model
-  * @param grad   gradient vector, which are just the linear coefficients
+  * @param expected computed from the model
+  * @param beta     gradient vector, which are just the linear coefficients
   */
-class LinearRegressionResult(values: Seq[Double], grad: Vector[Double]) extends PredictionResult[Double] {
-
-  /**
-    * Get the expected values for this prediction
-    *
-    * @return expected value of each prediction
-    */
-  override def getExpected(): Seq[Double] = values
+case class LinearRegressionPrediction(expected: Seq[Double], beta: Vector[Double]) extends PredictionResult[Double] {
 
   /**
     * Get the gradient, which is uniform
     *
     * @return a vector of doubles for each prediction
     */
-  override def getGradient(): Option[Seq[Vector[Double]]] = Some(Seq.fill(values.size)(grad))
+  override def gradient: Option[Seq[Vector[Double]]] = Some(Seq.fill(expected.size)(beta))
 }
diff --git a/src/main/scala/io/citrine/lolo/transformers/Standardizer.scala b/src/main/scala/io/citrine/lolo/transformers/Standardizer.scala
deleted file mode 100644
index 699cdbc1..00000000
--- a/src/main/scala/io/citrine/lolo/transformers/Standardizer.scala
+++ /dev/null
@@ -1,383 +0,0 @@
-package io.citrine.lolo.transformers
-
-import io.citrine.lolo._
-import io.citrine.random.Random
-
-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
-  *
-  * This is particularly helpful for regularized methods, like ridge regression, where
-  * the relative scale of the features and labels is important.
-  *
-  * Created by maxhutch on 2/19/17.
-  */
-case class Standardizer(baseLearner: Learner) extends Learner {
-
-  /**
-    * Create affine transformations for continuous features and labels; pass data through to learner
-    *
-    * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
-    * @param rng          random number generator for reproducibility
-    * @return training result containing a model
-    */
-  override def train(
-      trainingData: Seq[(Vector[Any], Any)],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): StandardizerTrainingResult = {
-    val inputTrans = Standardizer.getMultiStandardization(trainingData.map(_._1))
-    val outputTrans = trainingData.head._2 match {
-      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, rng)
-
-    new StandardizerTrainingResult(baseTrainingResult, outputTrans, inputTrans)
-  }
-}
-
-class MultiTaskStandardizer(baseLearner: MultiTaskLearner) extends MultiTaskLearner {
-
-  /**
-    * Train a model
-    *
-    * @param trainingData  to train on
-    * @param weights for the training rows, if applicable
-    * @param rng          random number generator for reproducibility
-    * @return a sequence of training results, one for each label
-    */
-  override def train(
-      trainingData: Seq[(Vector[Any], Vector[Any])],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): MultiTaskStandardizerTrainingResult = {
-    val (inputs, labels) = trainingData.unzip
-    val labelsTransposed = labels.transpose.toVector
-    val repOutput = labels.head
-    val inputTrans = Standardizer.getMultiStandardization(inputs)
-    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 = labelsTransposed
-      .zip(outputTrans)
-      .map {
-        case (labelSeq, trans) =>
-          Standardizer.applyStandardization(labelSeq, trans).toVector
-      }
-      .transpose
-
-    val baseTrainingResult = baseLearner.train(standardInputs.zip(standardLabels), weights, rng)
-    new MultiTaskStandardizerTrainingResult(baseTrainingResult, outputTrans, inputTrans)
-  }
-}
-
-/**
-  * 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 transformation (rescale, offset) of output label
-  * @param inputTrans sequence of optional transformations (rescale, offset) of inputs
-  */
-class StandardizerTrainingResult(
-    baseTrainingResult: TrainingResult,
-    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(), 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)
-            .lazyZip(Standardizer.invertStandardization(predictedVsActual.map(_._2), outputTrans))
-            .lazyZip(Standardizer.invertStandardization(predictedVsActual.map(_._3), outputTrans))
-            .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)
-            .lazyZip(Standardizer.invertStandardizationOption(predictedVsActual.map(_._2), outputTrans))
-            .lazyZip(Standardizer.invertStandardizationOption(predictedVsActual.map(_._3), outputTrans))
-            .toSeq
-        )
-    }
-  }
-}
-
-/**
-  * Model that wraps the base model next to the transformations
-  *
-  * @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]],
-    outputTrans: Option[Standardization],
-    inputTrans: Seq[Option[Standardization]]
-) extends Model[PredictionResult[T]] {
-
-  /**
-    * Standardize the inputs and then apply the base model
-    *
-    * @param inputs to apply the model to
-    * @return a predictionresult which includes, at least, the expected outputs
-    */
-  override def transform(inputs: Seq[Vector[Any]]): StandardizerPrediction[T] = {
-    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 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],
-    outputTrans: Option[Standardization],
-    inputTrans: Seq[Option[Standardization]]
-) extends PredictionResult[T] {
-
-  /**
-    * Get the expected values for this prediction
-    *
-    * Just reverse any transformation that was applied to the labels
-    *
-    * @return expected value of each prediction
-    */
-  override def getExpected(): Seq[T] = {
-    baseResult
-      .getExpected()
-      .map {
-        case x: Double if outputTrans.isDefined => outputTrans.get.invert(x)
-        case x: Any                             => x
-      }
-      .asInstanceOf[Seq[T]]
-  }
-
-  /**
-    * Get the uncertainty of the prediction
-    *
-    * This is un-standardized by rescaling by the label's variance
-    *
-    * @return uncertainty of each prediction
-    */
-  override def getUncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = {
-    baseResult.getUncertainty(includeNoise) match {
-      case Some(x) if outputTrans.isDefined => Some(x.map(_.asInstanceOf[Double] * outputRescale))
-      case x: Any                           => x
-    }
-  }
-
-  /**
-    * Get the gradient or sensitivity of each prediction
-    *
-    * This is un-stanardized by rescaling by the label's variance divided by the feature transformation's variance
-    *
-    * @return a vector of doubles for each prediction
-    */
-  override def getGradient(): Option[Seq[Vector[Double]]] = {
-    baseResult.getGradient() match {
-      case None => None
-      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 outputRescale = outputTrans.map(_.scale).getOrElse(1.0)
-}
-
-/**
-  * Utilities to compute and apply standarizations
-  */
-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 (shift, rescaling)
-    */
-  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 && scale < Double.PositiveInfinity) {
-      Standardization(mean, scale)
-    } else {
-      Standardization(mean, 1.0)
-    }
-  }
-
-  /**
-    * Get standardization for multiple values in a vector.
-    *
-    * This has a different name because the jvm erases the inner type of Seq[T]
-    *
-    * @param values sequence of vectors to be standardized
-    * @return sequence of standardization, each as an option
-    */
-  def getMultiStandardization(values: Seq[Vector[Any]]): Seq[Option[Standardization]] = {
-    val rep = values.head
-    rep.indices.map { i =>
-      rep(i) match {
-        case _: Double => Some(getStandardization(values.map(r => r(i).asInstanceOf[Double])))
-        case _: Any    => None
-      }
-    }
-  }
-
-  /**
-    * 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[Standardization]]): Seq[Vector[Any]] = {
-    input.map { r =>
-      r.zip(trans).map {
-        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
-      }
-    }
-  }
-
-  /**
-    * Apply the standardization to a sequence of values, which should result in output with zero mean and unit variance
-    *
-    * @param input to standardize
-    * @param trans transformation to apply
-    * @return sequence of standardized values
-    */
-  def applyStandardization(input: Seq[Any], trans: Option[Standardization]): Seq[Any] = {
-    if (trans.isEmpty) return input
-    input.asInstanceOf[Seq[Double]].map(trans.get.apply)
-  }
-
-  /**
-    * 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/transformers/FeatureRotator.scala b/src/main/scala/io/citrine/lolo/transformers/rotator/FeatureRotator.scala
similarity index 54%
rename from src/main/scala/io/citrine/lolo/transformers/FeatureRotator.scala
rename to src/main/scala/io/citrine/lolo/transformers/rotator/FeatureRotator.scala
index 97857eea..fa1f7cf1 100644
--- a/src/main/scala/io/citrine/lolo/transformers/FeatureRotator.scala
+++ b/src/main/scala/io/citrine/lolo/transformers/rotator/FeatureRotator.scala
@@ -1,4 +1,4 @@
-package io.citrine.lolo.transformers
+package io.citrine.lolo.transformers.rotator
 
 import io.citrine.lolo._
 import io.citrine.random.Random
@@ -6,131 +6,67 @@ import breeze.linalg.{diag, qr, DenseMatrix, DenseVector}
 import breeze.linalg.qr.QR
 import breeze.numerics.signum
 import breeze.stats.distributions.Gaussian
+import io.citrine.lolo.api.{Learner, Model, PredictionResult, TrainingResult, TrainingRow}
 import io.citrine.lolo.stats.StatsUtils.breezeRandBasis
 
 /**
-  * Rotate the training data before passing along to a base learner
+  * Rotate the training data before passing along to a base learner.
   *
   * This may be useful for improving randomization in random forests,
   * especially when using random feature selection without bagging.
   */
-case class FeatureRotator(baseLearner: Learner) extends Learner {
+case class FeatureRotator[T](baseLearner: Learner[T]) extends Learner[T] {
 
   /**
     * Create linear transformations for continuous features and labels & pass data through to learner
     *
     * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
     * @param rng          random number generator for reproducibility
     * @return training result containing a model
     */
-  override def train(
-      trainingData: Seq[(Vector[Any], Any)],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): RotatedFeatureTrainingResult = {
-    val featuresToRotate = FeatureRotator.getDoubleFeatures(trainingData.head._1)
+  override def train(trainingData: Seq[TrainingRow[T]], rng: Random): RotatedFeatureTrainingResult[T] = {
+    val featuresToRotate = FeatureRotator.getDoubleFeatures(trainingData.head.inputs)
     val trans = FeatureRotator.getRandomRotation(featuresToRotate.length, rng)
 
-    val (inputs, labels) = trainingData.unzip
-    val rotatedTrainingData = FeatureRotator.applyRotation(inputs, featuresToRotate, trans).zip(labels)
-    val baseTrainingResult = baseLearner.train(rotatedTrainingData, weights, rng)
+    val baseInputs = trainingData.map(_.inputs)
+    val rotatedInputs = FeatureRotator.applyRotation(baseInputs, featuresToRotate, trans)
+    val rotatedTrainingData = trainingData.zip(rotatedInputs).map {
+      case (row, rotated) => row.withInputs(rotated)
+    }
 
+    val baseTrainingResult = baseLearner.train(rotatedTrainingData, rng)
     RotatedFeatureTrainingResult(baseTrainingResult, featuresToRotate, trans)
   }
 }
 
-case class MultiTaskFeatureRotator(baseLearner: MultiTaskLearner) extends MultiTaskLearner {
-
-  /**
-    * Create linear transformations for continuous features and labels & pass data through to learner
-    *
-    * @param trainingData  to train on
-    * @param weights for the training rows, if applicable
-    * @param rng          random number generator for reproducibility
-    * @return a sequence of training results, one for each label
-    */
-  override def train(
-      trainingData: Seq[(Vector[Any], Vector[Any])],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): MultiTaskRotatedFeatureTrainingResult = {
-    val inputs = trainingData.map(_._1)
-    val labels = trainingData.map(_._2)
-    val featuresToRotate = FeatureRotator.getDoubleFeatures(inputs.head)
-    val trans = FeatureRotator.getRandomRotation(inputs.head.length, rng)
-    val rotatedFeatures = FeatureRotator.applyRotation(inputs, featuresToRotate, trans)
-    val baseTrainingResult = baseLearner.train(rotatedFeatures.zip(labels), weights, rng)
-    MultiTaskRotatedFeatureTrainingResult(baseTrainingResult, featuresToRotate, trans)
-  }
-}
-
 /**
-  * Training result bundling the base learner's training result with the list of rotated features and the transformation
+  * Training result bundling the base learner's training result with the list of rotated features and 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 RotatedFeatureTrainingResult(
-    baseTrainingResult: TrainingResult,
+case class RotatedFeatureTrainingResult[T](
+    baseTrainingResult: TrainingResult[T],
     rotatedFeatures: IndexedSeq[Int],
     trans: DenseMatrix[Double]
-) extends TrainingResult {
+) extends TrainingResult[T] {
 
-  /**
-    * Get the model contained in the training result
-    *
-    * @return the model
-    */
-  override def getModel(): Model[PredictionResult[Any]] = {
-    RotatedFeatureModel(baseTrainingResult.getModel(), rotatedFeatures, trans)
-  }
+  override def model: Model[T] = RotatedFeatureModel(baseTrainingResult.model, rotatedFeatures, trans)
 
-  override def getLoss(): Option[Double] = baseTrainingResult.getLoss()
+  override def loss: Option[Double] = baseTrainingResult.loss
 
-  override def getPredictedVsActual(): Option[Seq[(Vector[Any], Any, Any)]] = {
-    baseTrainingResult.getPredictedVsActual().map { x =>
+  override def predictedVsActual: Option[Seq[(Vector[Any], T, T)]] = {
+    baseTrainingResult.predictedVsActual.map { x =>
       x.map {
-        case (v: Vector[Any], e: Any, a: Any) => (FeatureRotator.applyOneRotation(v, rotatedFeatures, trans), e, a)
+        case (v, e, a) => (FeatureRotator.applyOneRotation(v, rotatedFeatures, trans), e, a)
       }
     }
   }
 }
 
 /**
-  * 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
+  * Model bundling the base learner's model with the list of rotated features and the transformation.
   *
   * @param baseModel model to which to delegate prediction on rotated features
   * @param rotatedFeatures indices of features to rotate
@@ -138,16 +74,16 @@ case class MultiTaskRotatedFeatureTrainingResult(
   * @tparam T label type
   */
 case class RotatedFeatureModel[T](
-    baseModel: Model[PredictionResult[T]],
+    baseModel: Model[T],
     rotatedFeatures: IndexedSeq[Int],
     trans: DenseMatrix[Double]
-) extends Model[PredictionResult[T]] {
+) extends Model[T] {
 
   /**
-    * Transform the inputs and then apply the base model
+    * Rotate the inputs and then apply the base model.
     *
     * @param inputs to apply the model to
-    * @return a RotatedFeaturePredictionResult which includes, at least, the expected outputs
+    * @return a RotatedFeaturePrediction which includes, at least, the expected outputs
     */
   override def transform(inputs: Seq[Vector[Any]]): RotatedFeaturePrediction[T] = {
     val rotatedInputs = FeatureRotator.applyRotation(inputs, rotatedFeatures, trans)
@@ -170,30 +106,24 @@ case class RotatedFeaturePrediction[T](
 ) extends PredictionResult[T] {
 
   /**
-    * Get the expected values for this prediction by delegating to baseResult
+    * Get the expected values for this prediction by delegating to the baseResult.
     *
     * @return expected value of each prediction
     */
-  override def getExpected(): Seq[T] = baseResult.getExpected()
+  override def expected: Seq[T] = baseResult.expected
 
   /**
-    * Get the uncertainty of the prediction by delegating to baseResult
+    * Get the uncertainty of the prediction by delegating to baseResult.
     *
     * @return uncertainty of each prediction
     */
-  override def getUncertainty(observational: Boolean): Option[Seq[Any]] = baseResult.getUncertainty(observational)
+  override def uncertainty(observational: Boolean): Option[Seq[Any]] = baseResult.uncertainty(observational)
 
-  /**
-    * Get the gradient or sensitivity of each prediction
-    *
-    * @return a vector of doubles for each prediction
-    */
-  override def getGradient(): Option[Seq[Vector[Double]]] = {
-    baseResult.getGradient().map { g =>
+  override def gradient: Option[Seq[Vector[Double]]] = {
+    baseResult.gradient.map { g =>
       FeatureRotator.applyRotation(g, rotatedFeatures, trans.t).asInstanceOf[Seq[Vector[Double]]]
     }
   }
-
 }
 
 /**
diff --git a/src/main/scala/io/citrine/lolo/transformers/rotator/MultiTaskFeatureRotator.scala b/src/main/scala/io/citrine/lolo/transformers/rotator/MultiTaskFeatureRotator.scala
new file mode 100644
index 00000000..cccf68be
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/transformers/rotator/MultiTaskFeatureRotator.scala
@@ -0,0 +1,67 @@
+package io.citrine.lolo.transformers.rotator
+
+import breeze.linalg.DenseMatrix
+import io.citrine.lolo.api.{
+  Model,
+  MultiTaskLearner,
+  MultiTaskModel,
+  MultiTaskTrainingResult,
+  ParallelModels,
+  TrainingRow
+}
+import io.citrine.random.Random
+
+case class MultiTaskFeatureRotator(baseLearner: MultiTaskLearner) extends MultiTaskLearner {
+
+  /**
+    * Create linear transformations for continuous features and labels & pass data through to learner
+    *
+    * @param trainingData  to train on
+    * @param rng          random number generator for reproducibility
+    * @return a sequence of training results, one for each label
+    */
+  override def train(
+      trainingData: Seq[TrainingRow[Vector[Any]]],
+      rng: Random
+  ): MultiTaskRotatedFeatureTrainingResult = {
+    val baseInputs = trainingData.map(_.inputs)
+    val featuresToRotate = FeatureRotator.getDoubleFeatures(baseInputs.head)
+    val trans = FeatureRotator.getRandomRotation(baseInputs.head.length, rng)
+
+    val rotatedInputs = FeatureRotator.applyRotation(baseInputs, featuresToRotate, trans)
+    val rotatedTrainingData = trainingData.zip(rotatedInputs).map {
+      case (row, rotated) => row.withInputs(rotated)
+    }
+
+    val baseTrainingResult = baseLearner.train(rotatedTrainingData, rng)
+    MultiTaskRotatedFeatureTrainingResult(baseTrainingResult, featuresToRotate, trans)
+  }
+}
+
+/**
+  * 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 model: MultiTaskModel = ParallelModels(models, baseTrainingResult.model.realLabels)
+
+  override def models: Seq[Model[Any]] =
+    baseTrainingResult.models.map { model => RotatedFeatureModel(model, rotatedFeatures, trans) }
+
+  override def predictedVsActual: Option[Seq[(Vector[Any], Vector[Option[Any]], Vector[Option[Any]])]] = {
+    baseTrainingResult.predictedVsActual.map { pva =>
+      pva.map {
+        case (inputs, predicted, actual) =>
+          (FeatureRotator.applyOneRotation(inputs, rotatedFeatures, trans), predicted, actual)
+      }
+    }
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/transformers/standardizer/MultiTaskStandardizer.scala b/src/main/scala/io/citrine/lolo/transformers/standardizer/MultiTaskStandardizer.scala
new file mode 100644
index 00000000..681d5188
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/transformers/standardizer/MultiTaskStandardizer.scala
@@ -0,0 +1,77 @@
+package io.citrine.lolo.transformers.standardizer
+
+import io.citrine.lolo.api.{
+  Model,
+  MultiTaskLearner,
+  MultiTaskModel,
+  MultiTaskTrainingResult,
+  ParallelModels,
+  TrainingRow
+}
+import io.citrine.random.Random
+
+/**
+  * Standardize the training data to zero mean and unit variance before feeding it into a multi-task learner.
+  *
+  * @param baseLearner base multi-task learner to train on standardized data.
+  */
+case class MultiTaskStandardizer(baseLearner: MultiTaskLearner) extends MultiTaskLearner {
+
+  override def train(trainingData: Seq[TrainingRow[Vector[Any]]], rng: Random): MultiTaskStandardizerTrainingResult = {
+    val (inputs, labels, weights) = trainingData.map(_.asTuple).unzip3
+
+    val inputTrans = Standardization.buildMulti(inputs)
+    val outputTrans = Standardization.buildMulti(labels)
+
+    val standardInputs = inputs.map { input => Standardization.applyMulti(input, inputTrans) }
+    val standardLabels = labels.map { label => Standardization.applyMulti(label, outputTrans) }
+    val standardTrainingData = standardInputs.lazyZip(standardLabels).lazyZip(weights).map {
+      case (i, l, w) => TrainingRow(i, l, w)
+    }
+
+    val baseTrainingResult = baseLearner.train(standardTrainingData, rng)
+    MultiTaskStandardizerTrainingResult(baseTrainingResult, outputTrans, inputTrans)
+  }
+}
+
+/**
+  * 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 sequence of optional transformation (rescale, offset) of output labels
+  * @param inputTrans sequence of optional transformations (rescale, offset) of inputs
+  */
+case class MultiTaskStandardizerTrainingResult(
+    baseTrainingResult: MultiTaskTrainingResult,
+    outputTrans: Seq[Option[Standardization]],
+    inputTrans: Seq[Option[Standardization]]
+) extends MultiTaskTrainingResult {
+
+  override def model: MultiTaskModel = ParallelModels(models, baseTrainingResult.model.realLabels)
+
+  override def models: Seq[StandardizerModel[Any]] = {
+    val realLabels = baseTrainingResult.model.realLabels
+    baseTrainingResult.models.zipWithIndex.map {
+      case (model, idx) =>
+        if (realLabels(idx)) {
+          RegressionStandardizerModel(model.asInstanceOf[Model[Double]], outputTrans(idx).get, inputTrans)
+        } else {
+          ClassificationStandardizerModel(model, inputTrans)
+        }
+    }
+  }
+
+  override def featureImportance: Option[Vector[Double]] = baseTrainingResult.featureImportance
+
+  override def predictedVsActual: Option[Seq[(Vector[Any], Vector[Option[Any]], Vector[Option[Any]])]] = {
+    baseTrainingResult.predictedVsActual.map { pva =>
+      pva.map {
+        case (inputs, predOpt, actualOpt) =>
+          val invertedInputs = Standardization.invertMulti(inputs, inputTrans)
+          val invertedPred = Standardization.invertMultiOption(predOpt, outputTrans)
+          val invertedActual = Standardization.invertMultiOption(actualOpt, outputTrans)
+          (invertedInputs, invertedPred, invertedActual)
+      }
+    }
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/transformers/standardizer/Standardization.scala b/src/main/scala/io/citrine/lolo/transformers/standardizer/Standardization.scala
new file mode 100644
index 00000000..3e5e5af2
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/transformers/standardizer/Standardization.scala
@@ -0,0 +1,108 @@
+package io.citrine.lolo.transformers.standardizer
+
+/** A package of the shift/scale required to standardize values to zero mean and unit variance. */
+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
+}
+
+object Standardization {
+
+  /**
+    * 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 (shift, rescaling)
+    */
+  def build(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 && scale < Double.PositiveInfinity) {
+      Standardization(mean, scale)
+    } else {
+      Standardization(mean, 1.0)
+    }
+  }
+
+  /**
+    * Get standardization for multiple values in a vector.
+    *
+    * This has a different name because the JVM erases the inner type of Seq[T].
+    *
+    * @param values sequence of vectors to be standardized
+    * @return sequence of standardization, each as an option
+    */
+  def buildMulti(values: Seq[Vector[Any]]): Seq[Option[Standardization]] = {
+    val rep = values.head
+    val valuesTransposed = values.transpose
+    rep.zipWithIndex.map {
+      case (_: Double, idx) =>
+        val doubleValues = valuesTransposed(idx).collect { case x: Double if !x.isNaN => x }
+        Some(Standardization.build(doubleValues))
+      case _ =>
+        None
+    }
+  }
+
+  /**
+    * Apply the standardizations to a vector, which should result in an output with zero mean and unit variance
+    *
+    * @param input to standardize
+    * @param trans transformations to apply
+    * @return a standardized vector
+    */
+  def applyMulti(input: Vector[Any], trans: Seq[Option[Standardization]]): Vector[Any] = {
+    input.zip(trans).map {
+      case (x: Double, Some(t)) => t.apply(x)
+      case (x, _)               => x
+    }
+  }
+
+  /**
+    * Apply the standardizations to a sequence of optional values.
+    *
+    * @param input to apply the standardization, if the value is defined
+    * @param trans transformations to un-apply
+    * @return a standardized vector of optional values
+    */
+  def applyMultiOption(input: Vector[Option[Any]], trans: Seq[Option[Standardization]]): Vector[Option[Any]] = {
+    input.zip(trans).map {
+      case (Some(x: Double), Some(t)) => Some(t.apply(x))
+      case (x, _)                     => x
+    }
+  }
+
+  /**
+    * Invert the standardization on a vector.
+    *
+    * @param input to invert the standardization
+    * @param trans transformations to un-apply
+    * @return a restored vector
+    */
+  def invertMulti(input: Vector[Any], trans: Seq[Option[Standardization]]): Vector[Any] = {
+    input.zip(trans).map {
+      case (x: Double, Some(t)) => t.invert(x)
+      case (x, _)               => x
+    }
+  }
+
+  /**
+    * Invert the standardizations on a sequence of optional values.
+    *
+    * @param input to invert the standardization, if the value is defined
+    * @param trans transformations to un-apply
+    * @return a restored vector of optional values
+    */
+  def invertMultiOption(input: Vector[Option[Any]], trans: Seq[Option[Standardization]]): Vector[Option[Any]] = {
+    input.zip(trans).map {
+      case (Some(x: Double), Some(t)) => Some(t.invert(x))
+      case (x, _)                     => x
+    }
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/transformers/standardizer/Standardizer.scala b/src/main/scala/io/citrine/lolo/transformers/standardizer/Standardizer.scala
new file mode 100644
index 00000000..59c161d4
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/transformers/standardizer/Standardizer.scala
@@ -0,0 +1,59 @@
+package io.citrine.lolo.transformers.standardizer
+
+import io.citrine.lolo.api.{Learner, TrainingRow}
+import io.citrine.random.Random
+
+trait Standardizer[T] extends Learner[T] {
+
+  /** The base learner trained on the standardized data. */
+  def baseLearner: Learner[T]
+
+  /**
+    * Create affine transformations for continuous features/labels; pass data through to the base learner.
+    *
+    * @param trainingData to train on
+    * @param rng          random number generator for reproducibility
+    * @return training result containing a model
+    */
+  override def train(trainingData: Seq[TrainingRow[T]], rng: Random = Random()): StandardizerTrainingResult[T]
+}
+
+/** A standardizer on both inputs & outputs for regression tasks. */
+case class RegressionStandardizer(baseLearner: Learner[Double]) extends Standardizer[Double] {
+
+  override def train(
+      trainingData: Seq[TrainingRow[Double]],
+      rng: Random = Random()
+  ): RegressionStandardizerTrainingResult = {
+    val (inputs, labels, weights) = trainingData.map(_.asTuple).unzip3
+
+    val inputTrans = Standardization.buildMulti(inputs)
+    val outputTrans = Standardization.build(labels)
+
+    val standardInputs = inputs.map { input => Standardization.applyMulti(input, inputTrans) }
+    val standardLabels = labels.map(outputTrans.apply)
+    val standardTrainingData = standardInputs.lazyZip(standardLabels).lazyZip(weights).map {
+      case (i, l, w) => TrainingRow(i, l, w)
+    }
+
+    val baseTrainingResult = baseLearner.train(standardTrainingData, rng)
+    RegressionStandardizerTrainingResult(baseTrainingResult, outputTrans, inputTrans)
+  }
+}
+
+/** A standardizer on features for classification tasks. */
+case class ClassificationStandardizer[T](baseLearner: Learner[T]) extends Standardizer[T] {
+
+  override def train(trainingData: Seq[TrainingRow[T]], rng: Random): ClassificationStandardizerTrainingResult[T] = {
+    val inputs = trainingData.map(_.inputs)
+    val inputTrans = Standardization.buildMulti(inputs)
+
+    val standardInputs = inputs.map { input => Standardization.applyMulti(input, inputTrans) }
+    val standardTrainingData = trainingData.zip(standardInputs).map {
+      case (row, inputs) => row.withInputs(inputs)
+    }
+
+    val baseTrainingResult = baseLearner.train(standardTrainingData, rng)
+    ClassificationStandardizerTrainingResult(baseTrainingResult, inputTrans)
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/transformers/standardizer/StandardizerModel.scala b/src/main/scala/io/citrine/lolo/transformers/standardizer/StandardizerModel.scala
new file mode 100644
index 00000000..87f1d1ac
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/transformers/standardizer/StandardizerModel.scala
@@ -0,0 +1,35 @@
+package io.citrine.lolo.transformers.standardizer
+
+import io.citrine.lolo.api.Model
+
+trait StandardizerModel[+T] extends Model[T] {
+
+  /** The base model trained on standardized data. */
+  def baseModel: Model[T]
+
+  /** Standardize the inputs and apply the base model. */
+  override def transform(inputs: Seq[Vector[Any]]): StandardizerPrediction[T]
+}
+
+case class RegressionStandardizerModel(
+    baseModel: Model[Double],
+    outputTrans: Standardization,
+    inputTrans: Seq[Option[Standardization]]
+) extends StandardizerModel[Double] {
+
+  override def transform(inputs: Seq[Vector[Any]]): RegressionStandardizerPrediction = {
+    val standardInputs = inputs.map { input => Standardization.applyMulti(input, inputTrans) }
+    RegressionStandardizerPrediction(baseModel.transform(standardInputs), outputTrans, inputTrans)
+  }
+}
+
+case class ClassificationStandardizerModel[T](
+    baseModel: Model[T],
+    inputTrans: Seq[Option[Standardization]]
+) extends StandardizerModel[T] {
+
+  override def transform(inputs: Seq[Vector[Any]]): ClassificationStandardizerPrediction[T] = {
+    val standardInputs = inputs.map { input => Standardization.applyMulti(input, inputTrans) }
+    ClassificationStandardizerPrediction(baseModel.transform(standardInputs), inputTrans)
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/transformers/standardizer/StandardizerPrediction.scala b/src/main/scala/io/citrine/lolo/transformers/standardizer/StandardizerPrediction.scala
new file mode 100644
index 00000000..76bdb5bb
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/transformers/standardizer/StandardizerPrediction.scala
@@ -0,0 +1,64 @@
+package io.citrine.lolo.transformers.standardizer
+
+import io.citrine.lolo.api.PredictionResult
+
+trait StandardizerPrediction[+T] extends PredictionResult[T] {
+
+  /** The base prediction from the model trained on standardized data. */
+  def basePrediction: PredictionResult[T]
+}
+
+case class RegressionStandardizerPrediction(
+    basePrediction: PredictionResult[Double],
+    outputTrans: Standardization,
+    inputTrans: Seq[Option[Standardization]]
+) extends StandardizerPrediction[Double] {
+
+  override def expected: Seq[Double] = basePrediction.expected.map(outputTrans.invert)
+
+  // TODO: A PredictionResult[Double] should always return a Option[Seq[Double]] for uncertainty
+  override def uncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = {
+    basePrediction.uncertainty(includeNoise).map { x =>
+      x.map(_.asInstanceOf[Double] * outputTrans.scale)
+    }
+  }
+
+  override def gradient: Option[Seq[Vector[Double]]] = {
+    basePrediction.gradient.map { gradients =>
+      gradients.map { g =>
+        g.zip(inputTrans).map {
+          case (y, inputStandardization) =>
+            // If there was a (linear) transformer used on that input, take the slope "m" and rescale by it
+            // Otherwise, just rescale by the output transformer
+            val inputScale = inputStandardization.map(_.scale).getOrElse(1.0)
+            val outputScale = outputTrans.scale
+            y * outputScale / inputScale
+        }
+      }
+    }
+  }
+}
+
+case class ClassificationStandardizerPrediction[T](
+    basePrediction: PredictionResult[T],
+    inputTrans: Seq[Option[Standardization]]
+) extends StandardizerPrediction[T] {
+
+  override def expected: Seq[T] = basePrediction.expected
+
+  override def uncertainty(includeNoise: Boolean = true): Option[Seq[Any]] =
+    basePrediction.uncertainty(includeNoise)
+
+  override def gradient: Option[Seq[Vector[Double]]] = {
+    basePrediction.gradient.map { gradients =>
+      gradients.map { g =>
+        g.zip(inputTrans).map {
+          case (y, inputStandardization) =>
+            // If there was a (linear) transformer used on that input, take the slope "m" and rescale by it
+            val inputScale = inputStandardization.map(_.scale).getOrElse(1.0)
+            y / inputScale
+        }
+      }
+    }
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/transformers/standardizer/StandardizerTrainingResult.scala b/src/main/scala/io/citrine/lolo/transformers/standardizer/StandardizerTrainingResult.scala
new file mode 100644
index 00000000..99ef404a
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/transformers/standardizer/StandardizerTrainingResult.scala
@@ -0,0 +1,62 @@
+package io.citrine.lolo.transformers.standardizer
+
+import io.citrine.lolo.api.TrainingResult
+
+/** Training result wrapping the base learner's training result next to the transformations. */
+trait StandardizerTrainingResult[+T] extends TrainingResult[T] {
+
+  def baseTrainingResult: TrainingResult[T]
+
+  override def model: StandardizerModel[T]
+
+  override def featureImportance: Option[Vector[Double]] = baseTrainingResult.featureImportance
+}
+
+/**
+  * Training result wrapping the input & output transformations for a regression learner.
+  *
+  * @param baseTrainingResult result of training on the standardized inputs and output
+  * @param outputTrans transformation (rescale, offset) of output label
+  * @param inputTrans sequence of optional transformations (rescale, offset) of inputs
+  */
+case class RegressionStandardizerTrainingResult(
+    baseTrainingResult: TrainingResult[Double],
+    outputTrans: Standardization,
+    inputTrans: Seq[Option[Standardization]]
+) extends StandardizerTrainingResult[Double] {
+
+  override def model: RegressionStandardizerModel =
+    RegressionStandardizerModel(baseTrainingResult.model, outputTrans, inputTrans)
+
+  override def predictedVsActual: Option[Seq[(Vector[Any], Double, Double)]] = {
+    baseTrainingResult.predictedVsActual.map { pva =>
+      pva.map {
+        case (inputs, pred, actual) =>
+          (Standardization.invertMulti(inputs, inputTrans), outputTrans.invert(pred), outputTrans.invert(actual))
+      }
+    }
+  }
+}
+
+/**
+  * Training result wrapping the input transformations for a classification learner.
+  *
+  * @param baseTrainingResult result of training on the standardized inputs and output
+  * @param inputTrans sequence of optional transformations (rescale, offset) of inputs
+  */
+case class ClassificationStandardizerTrainingResult[T](
+    baseTrainingResult: TrainingResult[T],
+    inputTrans: Seq[Option[Standardization]]
+) extends StandardizerTrainingResult[T] {
+
+  override def model: ClassificationStandardizerModel[T] =
+    ClassificationStandardizerModel(baseTrainingResult.model, inputTrans)
+
+  override def predictedVsActual: Option[Seq[(Vector[Any], T, T)]] = {
+    baseTrainingResult.predictedVsActual.map { pva =>
+      pva.map {
+        case (inputs, pred, actual) => (Standardization.invertMulti(inputs, inputTrans), pred, actual)
+      }
+    }
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/trees/Nodes.scala b/src/main/scala/io/citrine/lolo/trees/ModelNode.scala
similarity index 81%
rename from src/main/scala/io/citrine/lolo/trees/Nodes.scala
rename to src/main/scala/io/citrine/lolo/trees/ModelNode.scala
index 67b8ee92..739cc62b 100644
--- a/src/main/scala/io/citrine/lolo/trees/Nodes.scala
+++ b/src/main/scala/io/citrine/lolo/trees/ModelNode.scala
@@ -1,38 +1,13 @@
 package io.citrine.lolo.trees
 
 import breeze.linalg.DenseMatrix
+import io.citrine.lolo.api.{Model, PredictionResult}
 import io.citrine.lolo.trees.splits.Split
-import io.citrine.lolo.{Model, PredictionResult, TrainingResult}
 
-import scala.collection.mutable
+trait ModelNode[+T] extends Serializable {
 
-/**
-  * Trait to provide getNode interface for internal and leaf training nodes
-  *
-  * @tparam T type of the model output
-  */
-trait TrainingNode[+T] extends Serializable {
-
-  // TODO (PLA-10433): make this a structured type
-  def trainingData: Seq[(Vector[AnyVal], T, Double)]
-
-  /**
-    * Get the lightweight prediction node for the output tree
-    *
-    * @return lightweight prediction node
-    */
-  def modelNode: ModelNode[PredictionResult[T]]
-
-  /**
-    * Get the feature importance of the subtree below this node
-    *
-    * @return feature importance as a vector
-    */
-  def featureImportance: mutable.ArraySeq[Double]
-}
-
-trait ModelNode[+T <: PredictionResult[Any]] extends Serializable {
-  def transform(input: Vector[AnyVal]): (T, TreeMeta)
+  /** Transform the inputs, returning a prediction result and meta information regarding the tree depth. */
+  def transform(input: Vector[AnyVal]): (PredictionResult[T], TreeMeta)
 
   /**
     * Compute Shapley feature attributions for a given input in this node's subtree
@@ -70,7 +45,7 @@ trait ModelNode[+T <: PredictionResult[Any]] extends Serializable {
     *
     * @return total weight of training weight in subtree
     */
-  private[lolo] def getTrainingWeight(): Double
+  private[lolo] def trainingWeight: Double
 }
 
 /**
@@ -84,7 +59,7 @@ trait ModelNode[+T <: PredictionResult[Any]] extends Serializable {
   * @param trainingWeight  weight of training data in subtree (i.e. size of unweighted training set)
   * @tparam T type of the output
   */
-case class InternalModelNode[T <: PredictionResult[Any]](
+case class InternalModelNode[+T](
     split: Split,
     left: ModelNode[T],
     right: ModelNode[T],
@@ -98,7 +73,7 @@ case class InternalModelNode[T <: PredictionResult[Any]](
     * @param input to predict for
     * @return prediction
     */
-  override def transform(input: Vector[AnyVal]): (T, TreeMeta) = {
+  override def transform(input: Vector[AnyVal]): (PredictionResult[T], TreeMeta) = {
     if (split.turnLeft(input)) {
       left.transform(input)
     } else {
@@ -135,8 +110,8 @@ case class InternalModelNode[T <: PredictionResult[Any]](
 
     // If the feature is omitted or unknown, then the weight assigned to each child is taken
     // as the share of the training data that got split into each child
-    val hotPortion = hot.getTrainingWeight() / trainingWeight
-    val coldPortion = cold.getTrainingWeight() / trainingWeight
+    val hotPortion = hot.trainingWeight / trainingWeight
+    val coldPortion = cold.trainingWeight / trainingWeight
 
     if (omitFeatures.contains(featureIndex)) {
       // Don't add the feature to the featureWeights
@@ -173,28 +148,12 @@ case class InternalModelNode[T <: PredictionResult[Any]](
       coldContrib + hotContrib
     }
   }
-
-  override def getTrainingWeight(): Double = trainingWeight
 }
 
-/** A leaf defined by a training result. */
-trait TrainingLeaf[T] extends TrainingNode[T] {
-
-  def trainingResult: TrainingResult
-
-  def depth: Int
+case class ModelLeaf[+T](model: Model[T], depth: Int, trainingWeight: Double) extends ModelNode[T] {
 
-  def modelNode: ModelNode[PredictionResult[T]] =
-    ModelLeaf(model, depth, trainingData.map(_._3).sum)
-
-  def model: Model[PredictionResult[T]] = trainingResult.getModel().asInstanceOf[Model[PredictionResult[T]]]
-}
-
-case class ModelLeaf[T](model: Model[PredictionResult[T]], depth: Int, trainingWeight: Double)
-    extends ModelNode[PredictionResult[T]] {
-  override def transform(input: Vector[AnyVal]): (PredictionResult[T], TreeMeta) = {
+  override def transform(input: Vector[AnyVal]): (PredictionResult[T], TreeMeta) =
     (model.transform(Seq(input)), TreeMeta(depth))
-  }
 
   /**
     * Compute the contribution to SHAP in the leaf based on the features that were encountered between the root node
@@ -216,7 +175,7 @@ case class ModelLeaf[T](model: Model[PredictionResult[T]], depth: Int, trainingW
     featureWeights.values.foreach { case FeatureWeightFactor(exclude, include) => set.extend(exclude, include) }
 
     // The contribution is proportional to the leaf's prediction, so grab that
-    this.model.transform(Seq(input)).getExpected().head match {
+    this.model.transform(Seq(input)).expected.head match {
       case v: Double =>
         // For each feature, compute the contribution and store it in shapValues
         featureWeights.foreach {
@@ -232,7 +191,5 @@ case class ModelLeaf[T](model: Model[PredictionResult[T]], depth: Int, trainingW
     shapValues
   }
 
-  override def getTrainingWeight(): Double = trainingWeight
-
   override def shapley(input: Vector[AnyVal], omitFeatures: Set[Int] = Set()): Option[DenseMatrix[Double]] = None
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/TrainingNode.scala b/src/main/scala/io/citrine/lolo/trees/TrainingNode.scala
new file mode 100644
index 00000000..43603dbf
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/trees/TrainingNode.scala
@@ -0,0 +1,41 @@
+package io.citrine.lolo.trees
+
+import io.citrine.lolo.api.{Model, TrainingResult, TrainingRow}
+
+import scala.collection.mutable
+
+/**
+  * Trait to provide getNode interface for internal and leaf training nodes
+  *
+  * @tparam T type of the model output
+  */
+trait TrainingNode[+T] extends Serializable {
+
+  def trainingData: Seq[TrainingRow[T]]
+
+  /**
+    * Get the lightweight prediction node for the output tree
+    *
+    * @return lightweight prediction node
+    */
+  def modelNode: ModelNode[T]
+
+  /**
+    * Get the feature importance of the subtree below this node
+    *
+    * @return feature importance as a vector
+    */
+  def featureImportance: mutable.ArraySeq[Double]
+}
+
+/** A leaf defined by a training result. */
+trait TrainingLeaf[+T] extends TrainingNode[T] {
+
+  def depth: Int
+
+  def trainingResult: TrainingResult[T]
+
+  def modelNode: ModelNode[T] = ModelLeaf(model, depth, trainingData.map(_.weight).sum)
+
+  def model: Model[T] = trainingResult.model
+}
diff --git a/src/main/scala/io/citrine/lolo/trees/TreeMeta.scala b/src/main/scala/io/citrine/lolo/trees/TreeMeta.scala
index 48da3f99..9f75d40e 100644
--- a/src/main/scala/io/citrine/lolo/trees/TreeMeta.scala
+++ b/src/main/scala/io/citrine/lolo/trees/TreeMeta.scala
@@ -1,6 +1,4 @@
 package io.citrine.lolo.trees
 
-/**
-  * Created by maxhutch on 12/21/16.
-  */
+/** Meta-information about the current depth within a tree. */
 case class TreeMeta(depth: Int)
diff --git a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingLeaf.scala b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingLeaf.scala
index 01a5b321..6e53b711 100644
--- a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingLeaf.scala
+++ b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingLeaf.scala
@@ -1,6 +1,6 @@
 package io.citrine.lolo.trees.classification
 
-import io.citrine.lolo.{Learner, TrainingResult}
+import io.citrine.lolo.api.{Learner, TrainingResult, TrainingRow}
 import io.citrine.lolo.trees.TrainingLeaf
 import io.citrine.random.Random
 
@@ -8,12 +8,12 @@ import scala.collection.mutable
 
 /** Training leaf for a classification tree. */
 case class ClassificationTrainingLeaf(
-    trainingData: Seq[(Vector[AnyVal], Char, Double)],
-    trainingResult: TrainingResult,
+    trainingData: Seq[TrainingRow[Char]],
+    trainingResult: TrainingResult[Char],
     depth: Int
 ) extends TrainingLeaf[Char] {
 
-  override def featureImportance: mutable.ArraySeq[Double] = mutable.ArraySeq.fill(trainingData.head._1.size)(0.0)
+  override def featureImportance: mutable.ArraySeq[Double] = mutable.ArraySeq.fill(trainingData.head.inputs.size)(0.0)
 }
 
 object ClassificationTrainingLeaf {
@@ -26,8 +26,8 @@ object ClassificationTrainingLeaf {
     * @return a trained classification leaf node
     */
   def build(
-      trainingData: Seq[(Vector[AnyVal], Char, Double)],
-      leafLearner: Learner,
+      trainingData: Seq[TrainingRow[Char]],
+      leafLearner: Learner[Char],
       depth: Int,
       rng: Random
   ): ClassificationTrainingLeaf =
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 5f49e1b5..3dad25aa 100644
--- a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingNode.scala
+++ b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingNode.scala
@@ -2,14 +2,16 @@ package io.citrine.lolo.trees.classification
 
 import io.citrine.lolo.trees.splits.{NoSplit, Split, Splitter}
 import io.citrine.lolo.trees.{InternalModelNode, ModelNode, TrainingNode}
-import io.citrine.lolo.{Learner, PredictionResult}
+import io.citrine.lolo.api.{Learner, TrainingRow}
 import io.citrine.random.Random
 
+import scala.collection.mutable
+
 case class ClassificationTrainingNode(
-    trainingData: Seq[(Vector[AnyVal], Char, Double)],
+    trainingData: Seq[TrainingRow[Char]],
     leftNode: TrainingNode[Char],
     rightNode: TrainingNode[Char],
-    leafLearner: Learner,
+    leafLearner: Learner[Char],
     split: Split,
     deltaImpurity: Double,
     numClasses: Int
@@ -20,16 +22,16 @@ case class ClassificationTrainingNode(
     *
     * @return lightweight prediction node
     */
-  override def modelNode: ModelNode[PredictionResult[Char]] =
+  override def modelNode: ModelNode[Char] =
     InternalModelNode(
       split,
       leftNode.modelNode,
       rightNode.modelNode,
       numClasses,
-      trainingData.map(_._3).sum
+      trainingData.map(_.weight).sum
     )
 
-  override def featureImportance: scala.collection.mutable.ArraySeq[Double] = {
+  override def featureImportance: mutable.ArraySeq[Double] = {
     val improvement = deltaImpurity
     val ans = leftNode.featureImportance.zip(rightNode.featureImportance).map(p => p._1 + p._2)
     ans(split.index) = ans(split.index) + improvement
@@ -40,8 +42,8 @@ case class ClassificationTrainingNode(
 object ClassificationTrainingNode {
 
   def build(
-      trainingData: Seq[(Vector[AnyVal], Char, Double)],
-      leafLearner: Learner,
+      trainingData: Seq[TrainingRow[Char]],
+      leafLearner: Learner[Char],
       splitter: Splitter[Char],
       numFeatures: Int,
       minLeafInstances: Int,
@@ -52,7 +54,7 @@ object ClassificationTrainingNode {
   ): TrainingNode[Char] = {
     val sufficientData = trainingData.size >= 2 * minLeafInstances &&
       remainingDepth > 0 &&
-      trainingData.exists(_._2 != trainingData.head._2)
+      trainingData.exists(_.label != trainingData.head.label)
     val (split: Split, deltaImpurity: Double) = if (sufficientData) {
       splitter.getBestSplit(trainingData, numFeatures, minLeafInstances, rng)
     } else {
@@ -62,7 +64,7 @@ object ClassificationTrainingNode {
       case _: NoSplit =>
         ClassificationTrainingLeaf.build(trainingData, leafLearner, maxDepth - remainingDepth, rng)
       case split: Split =>
-        val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r._1))
+        val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r.inputs))
         val leftNode = ClassificationTrainingNode.build(
           trainingData = leftTrain,
           leafLearner = leafLearner,
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 6701ca89..9b53eb3f 100644
--- a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTree.scala
+++ b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTree.scala
@@ -1,11 +1,13 @@
 package io.citrine.lolo.trees.classification
 
+import io.citrine.lolo.api.{Learner, Model, PredictionResult, TrainingResult, TrainingRow}
 import io.citrine.random.Random
 import io.citrine.lolo.encoders.CategoricalEncoder
-import io.citrine.lolo.linear.GuessTheMeanLearner
-import io.citrine.lolo.trees.splits.{ClassificationSplitter, NoSplit, Splitter}
-import io.citrine.lolo.trees.{ModelNode, TrainingLeaf, TrainingNode, TreeMeta}
-import io.citrine.lolo.{Learner, Model, PredictionResult, TrainingResult}
+import io.citrine.lolo.linear.GuessTheModeLearner
+import io.citrine.lolo.trees.splits.{ClassificationSplitter, Splitter}
+import io.citrine.lolo.trees.{ModelNode, TrainingNode, TreeMeta}
+
+import scala.collection.mutable
 
 /**
   * @param numFeatures subset of features to select splits from
@@ -18,59 +20,50 @@ case class ClassificationTreeLearner(
     numFeatures: Int = -1,
     maxDepth: Int = 30,
     minLeafInstances: Int = 1,
-    leafLearner: Option[Learner] = None,
+    leafLearner: Option[Learner[Char]] = None,
     splitter: Splitter[Char] = ClassificationSplitter()
-) extends Learner {
+) extends Learner[Any] {
 
-  @transient private lazy val myLeafLearner: Learner = leafLearner.getOrElse(GuessTheMeanLearner())
+  @transient private lazy val myLeafLearner: Learner[Char] = leafLearner.getOrElse(GuessTheModeLearner())
 
   /**
-    * Train classification tree
+    * Train a classification tree.
     *
     * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
     * @param rng          random number generator for reproducibility
     * @return a classification tree
     */
-  override def train(
-      trainingData: Seq[(Vector[Any], Any)],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): ClassificationTrainingResult = {
+  override def train(trainingData: Seq[TrainingRow[Any]], rng: Random): ClassificationTreeTrainingResult = {
     assert(trainingData.size > 4, s"We need to have at least 4 rows, only ${trainingData.size} given")
-    val repInput = trainingData.head._1
+    val repInput = trainingData.head.inputs
 
-    /* Create encoders for any categorical features */
+    // Create encoders for any categorical features
     val inputEncoders: Seq[Option[CategoricalEncoder[Any]]] = repInput.zipWithIndex.map {
       case (v, i) =>
         if (v.isInstanceOf[Double]) {
           None
         } else {
-          Some(CategoricalEncoder.buildEncoder(trainingData.map(_._1(i))))
+          Some(CategoricalEncoder.buildEncoder(trainingData.map(_.inputs(i))))
         }
     }
 
-    val outputEncoder = CategoricalEncoder.buildEncoder(trainingData.map(_._2))
+    val outputEncoder = CategoricalEncoder.buildEncoder(trainingData.map(_.label))
 
-    /* Encode the training data */
-    val encodedTraining =
-      trainingData.map(p => (CategoricalEncoder.encodeInput(p._1, inputEncoders), outputEncoder.encode(p._2)))
+    // Encode the training data
+    val encodedTraining = trainingData.map { row =>
+      val encodedInputs = CategoricalEncoder.encodeInput(row.inputs, inputEncoders)
+      val encodedLabels = outputEncoder.encode(row.label)
+      TrainingRow(encodedInputs, encodedLabels, row.weight)
+    }
 
     /* Add the weights to the (features, label) tuples and remove any with zero weight */
-    val finalTraining = encodedTraining
-      .zip(weights.getOrElse(Seq.fill(trainingData.size)(1.0)))
-      .map {
-        case ((f, l), w) =>
-          (f, l, w)
-      }
-      .filter(_._3 > 0)
-      .toVector
+    val finalTraining = encodedTraining.filter(_.weight > 0.0)
 
     /* If the number of features isn't specified, use all of them */
     val numFeaturesActual = if (numFeatures > 0) {
       numFeatures
     } else {
-      finalTraining.head._1.size
+      finalTraining.head.inputs.size
     }
 
     // Recursively build the tree via its nodes and wrap the top node in a ClassificationTreeTrainingResult
@@ -82,49 +75,43 @@ case class ClassificationTreeLearner(
       minLeafInstances = minLeafInstances,
       remainingDepth = maxDepth,
       maxDepth = maxDepth,
-      numClasses = trainingData.map(_._2).distinct.length,
+      numClasses = trainingData.map(_.label).distinct.length,
       rng = rng
     )
-    new ClassificationTrainingResult(rootTrainingNode, inputEncoders, outputEncoder)
+    ClassificationTreeTrainingResult(rootTrainingNode, inputEncoders, outputEncoder)
   }
 }
 
-class ClassificationTrainingResult(
+case class ClassificationTreeTrainingResult(
     rootTrainingNode: TrainingNode[Char],
     inputEncoders: Seq[Option[CategoricalEncoder[Any]]],
     outputEncoder: CategoricalEncoder[Any]
-) extends TrainingResult {
-  /* Grab a prediction node.  The partitioning happens here */
-  lazy val model = new ClassificationTree(rootTrainingNode.modelNode, inputEncoders, outputEncoder)
-
-  /* Grab the feature influences */
-  lazy val importance = rootTrainingNode.featureImportance
-  private lazy val importanceNormalized = {
-    if (Math.abs(importance.sum) > 0) {
-      importance.map(_ / importance.sum)
-    } else {
-      importance.map(_ => 1.0 / importance.size)
-    }
-  }
+) extends TrainingResult[Any] {
 
-  override def getModel(): ClassificationTree = model
+  // Grab a prediction node. The partitioning happens here
+  override lazy val model: ClassificationTree =
+    ClassificationTree(rootTrainingNode.modelNode, inputEncoders, outputEncoder)
 
-  /**
-    * Get a measure of the importance of the model features
-    *
-    * @return feature influences as an array of doubles
-    */
-  override def getFeatureImportance(): Option[Vector[Double]] = Some(importanceNormalized.toVector)
+  // Grab the feature influences
+  lazy val nodeImportance: mutable.ArraySeq[Double] = rootTrainingNode.featureImportance
+
+  override lazy val featureImportance: Option[Vector[Double]] = Some(
+    if (Math.abs(nodeImportance.sum) > 0) {
+      nodeImportance.map(_ / nodeImportance.sum).toVector
+    } else {
+      nodeImportance.map(_ => 1.0 / nodeImportance.size).toVector
+    }
+  )
 }
 
 /**
   * Classification tree
   */
-class ClassificationTree(
-    rootModelNode: ModelNode[PredictionResult[Char]],
+case class ClassificationTree(
+    rootModelNode: ModelNode[Char],
     inputEncoders: Seq[Option[CategoricalEncoder[Any]]],
     outputEncoder: CategoricalEncoder[Any]
-) extends Model[ClassificationResult] {
+) extends Model[Any] {
 
   /**
     * Apply the model to a seq of inputs
@@ -132,8 +119,8 @@ class ClassificationTree(
     * @param inputs to apply the model to
     * @return a predictionresult which includes, at least, the expected outputs
     */
-  override def transform(inputs: Seq[Vector[Any]]): ClassificationResult = {
-    new ClassificationResult(
+  override def transform(inputs: Seq[Vector[Any]]): ClassificationTreePrediction = {
+    ClassificationTreePrediction(
       inputs.map(inp => rootModelNode.transform(CategoricalEncoder.encodeInput(inp, inputEncoders))),
       outputEncoder
     )
@@ -143,7 +130,7 @@ class ClassificationTree(
 /**
   * Classification result
   */
-class ClassificationResult(
+case class ClassificationTreePrediction(
     predictions: Seq[(PredictionResult[Char], TreeMeta)],
     outputEncoder: CategoricalEncoder[Any]
 ) extends PredictionResult[Any] {
@@ -153,9 +140,7 @@ class ClassificationResult(
     *
     * @return expected value of each prediction
     */
-  override def getExpected(): Seq[Any] = predictions.map(p => outputEncoder.decode(p._1.getExpected().head))
+  override def expected: Seq[Any] = predictions.map(p => outputEncoder.decode(p._1.expected.head))
 
-  def getDepth(): Seq[Int] = {
-    predictions.map(_._2.depth)
-  }
+  def depth: Seq[Int] = predictions.map(_._2.depth)
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/impurity/GiniCalculator.scala b/src/main/scala/io/citrine/lolo/trees/impurity/GiniCalculator.scala
index 5fef2972..4b121f3a 100644
--- a/src/main/scala/io/citrine/lolo/trees/impurity/GiniCalculator.scala
+++ b/src/main/scala/io/citrine/lolo/trees/impurity/GiniCalculator.scala
@@ -62,7 +62,7 @@ class GiniCalculator(
     }
   }
 
-  private val leftCategoryWeights = new Array[Double](totalCategoryWeights.size)
+  private val leftCategoryWeights = new Array[Double](totalCategoryWeights.length)
   private var leftWeight: Double = 0.0
   private var leftSquareSum: Double = 0.0
   private var rightSquareSum: Double = totalSquareSum
diff --git a/src/main/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculator.scala b/src/main/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculator.scala
index 250926b3..2dd2b31f 100644
--- a/src/main/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculator.scala
+++ b/src/main/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculator.scala
@@ -7,7 +7,7 @@ package io.citrine.lolo.trees.impurity
   */
 class MultiImpurityCalculator(
     calculators: Seq[ImpurityCalculator[AnyVal]]
-) extends ImpurityCalculator[Seq[AnyVal]] {
+) extends ImpurityCalculator[Vector[AnyVal]] {
 
   /**
     * Add the value to each calculator
@@ -16,13 +16,15 @@ class MultiImpurityCalculator(
     * @param weight of the value
     * @return the impurity after adding
     */
-  def add(value: Seq[AnyVal], weight: Double): Double = {
+  def add(value: Vector[AnyVal], weight: Double): Double = {
     value.zip(calculators).map {
       case (v, calc) =>
-        if (v.isInstanceOf[Double]) {
-          calc.asInstanceOf[ImpurityCalculator[Double]].add(v.asInstanceOf[Double], weight)
-        } else if (v.isInstanceOf[Char]) {
-          calc.asInstanceOf[ImpurityCalculator[Char]].add(v.asInstanceOf[Char], weight)
+        v match {
+          case d: Double =>
+            calc.asInstanceOf[ImpurityCalculator[Double]].add(d, weight)
+          case c: Char =>
+            calc.asInstanceOf[ImpurityCalculator[Char]].add(c, weight)
+          case _ =>
         }
     }
 
@@ -36,13 +38,15 @@ class MultiImpurityCalculator(
     * @param weight of the value
     * @return the impurity after removing
     */
-  def remove(value: Seq[AnyVal], weight: Double): Double = {
+  def remove(value: Vector[AnyVal], weight: Double): Double = {
     value.zip(calculators).map {
       case (v, calc) =>
-        if (v.isInstanceOf[Double]) {
-          calc.asInstanceOf[ImpurityCalculator[Double]].remove(v.asInstanceOf[Double], weight)
-        } else if (v.isInstanceOf[Char]) {
-          calc.asInstanceOf[ImpurityCalculator[Char]].remove(v.asInstanceOf[Char], weight)
+        v match {
+          case d: Double =>
+            calc.asInstanceOf[ImpurityCalculator[Double]].remove(d, weight)
+          case c: Char =>
+            calc.asInstanceOf[ImpurityCalculator[Char]].remove(c, weight)
+          case _ =>
         }
     }
 
@@ -78,7 +82,7 @@ object MultiImpurityCalculator {
     * @param weights which are assumed to be constant over the labels at each row
     * @return MultiImpurityCalculator that sum the impurity of each label index
     */
-  def build(labels: Seq[Seq[AnyVal]], weights: Seq[Double]): MultiImpurityCalculator = {
+  def build(labels: Seq[Vector[AnyVal]], weights: Seq[Double]): MultiImpurityCalculator = {
     val calculators: Seq[ImpurityCalculator[AnyVal]] = labels.transpose.map { labelSeq =>
       if (labelSeq.head.isInstanceOf[Double]) {
         VarianceCalculator
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 0c43e323..7ccc2d67 100644
--- a/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTrainingNode.scala
+++ b/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTrainingNode.scala
@@ -1,7 +1,7 @@
 package io.citrine.lolo.trees.multitask
 
-import io.citrine.lolo.PredictionResult
-import io.citrine.lolo.linear.GuessTheMeanLearner
+import io.citrine.lolo.api.{PredictionResult, TrainingRow}
+import io.citrine.lolo.linear.{GuessTheMeanLearner, GuessTheModeLearner}
 import io.citrine.lolo.trees.classification.ClassificationTrainingLeaf
 import io.citrine.lolo.trees.regression.RegressionTrainingLeaf
 import io.citrine.lolo.trees.splits.{MultiTaskSplitter, NoSplit, Split}
@@ -11,7 +11,7 @@ import io.citrine.random.Random
 import scala.collection.mutable
 
 case class MultiTaskTrainingNode(
-    trainingData: Seq[(Vector[AnyVal], Vector[AnyVal], Double)],
+    trainingData: Seq[TrainingRow[Vector[AnyVal]]],
     labelWiseInstructions: Seq[MultiTaskLabelInstruction],
     deltaImpurity: Double
 ) {
@@ -33,18 +33,18 @@ case class MultiTaskTrainingNode(
   }
 
   // Construct the model node for the `index`th label
-  def modelNodeByLabelIndex(index: Int): ModelNode[PredictionResult[Any]] = {
+  def modelNodeByLabelIndex(index: Int): ModelNode[Any] = {
     labelWiseInstructions(index) match {
       case Stop(leaf)             => leaf.modelNode
       case FollowChild(childNode) => childNode.modelNodeByLabelIndex(index)
       case DoSplit(split, leftNode, rightNode) =>
-        if (trainingData.head._2(index).isInstanceOf[Double]) {
+        if (trainingData.head.label(index).isInstanceOf[Double]) {
           new InternalModelNode[PredictionResult[Double]](
             split = split,
             left = leftNode.modelNodeByLabelIndex(index).asInstanceOf[ModelNode[PredictionResult[Double]]],
             right = rightNode.modelNodeByLabelIndex(index).asInstanceOf[ModelNode[PredictionResult[Double]]],
             outputDimension = 0, // Shapley is not implemented for multi-task nodes
-            trainingWeight = trainingData.map(_._3).sum
+            trainingWeight = trainingData.map(_.weight).sum
           )
         } else {
           new InternalModelNode[PredictionResult[Char]](
@@ -52,7 +52,7 @@ case class MultiTaskTrainingNode(
             left = leftNode.modelNodeByLabelIndex(index).asInstanceOf[ModelNode[PredictionResult[Char]]],
             right = rightNode.modelNodeByLabelIndex(index).asInstanceOf[ModelNode[PredictionResult[Char]]],
             outputDimension = 0,
-            trainingWeight = trainingData.map(_._3).sum
+            trainingWeight = trainingData.map(_.weight).sum
           )
         }
       case Inaccessible() => throw new RuntimeException(s"No valid training data present for label $index")
@@ -75,7 +75,7 @@ object MultiTaskTrainingNode {
     * @return
     */
   def build(
-      trainingData: Seq[(Vector[AnyVal], Vector[AnyVal], Double)],
+      trainingData: Seq[TrainingRow[Vector[AnyVal]]],
       numFeatures: Int,
       remainingDepth: Int,
       maxDepth: Int,
@@ -86,7 +86,7 @@ object MultiTaskTrainingNode {
     // Determine the overall split.
     val sufficientData = trainingData.size >= 2 * minInstances &&
       remainingDepth > 0 &&
-      trainingData.exists(row => row._2 != trainingData.head._2)
+      trainingData.exists(row => row.label != trainingData.head.label)
     val (split: Split, deltaImpurity: Double) = if (sufficientData) {
       splitter.getBestSplit(
         trainingData,
@@ -101,7 +101,7 @@ object MultiTaskTrainingNode {
     val nodesOpt: Option[(MultiTaskTrainingNode, MultiTaskTrainingNode)] = split match {
       case _: NoSplit => None
       case split: Split =>
-        val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r._1))
+        val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r.inputs))
         val leftNode = MultiTaskTrainingNode.build(
           trainingData = leftTrain,
           numFeatures = numFeatures,
@@ -124,15 +124,15 @@ object MultiTaskTrainingNode {
     }
     // Determine what to do for each label (evaluate a model, evaluate the split, go left, or go right).
     val exampleRow = trainingData.head
-    val labelWiseInstructions = exampleRow._2.indices.map { index =>
+    val labelWiseInstructions = exampleRow.label.indices.map { index =>
       // Determine how much data *with this label* goes down each branch
-      val exampleLabel = exampleRow._2(index)
-      val reducedData = if (exampleLabel.isInstanceOf[Double]) {
-        trainingData.map(x => (x._1, x._2(index).asInstanceOf[Double], x._3)).filterNot(_._2.isNaN)
+      val exampleLabel = exampleRow.label(index)
+      val reducedData: Seq[TrainingRow[Any]] = if (exampleLabel.isInstanceOf[Double]) {
+        trainingData.map(_.mapLabel(labelVec => labelVec(index).asInstanceOf[Double])).filterNot(_.label.isNaN)
       } else {
-        trainingData.map(x => (x._1, x._2(index).asInstanceOf[Char], x._3)).filter(_._2 > 0)
+        trainingData.map(_.mapLabel(labelVec => labelVec(index).asInstanceOf[Char])).filter(_.label > 0)
       }
-      val (left, right) = reducedData.partition(r => split.turnLeft(r._1))
+      val (left, right) = reducedData.partition(r => split.turnLeft(r.inputs))
 
       if (reducedData.isEmpty) {
         Inaccessible()
@@ -141,7 +141,7 @@ object MultiTaskTrainingNode {
         val trainingLeaf: TrainingNode[AnyVal] = if (exampleLabel.isInstanceOf[Double]) {
           RegressionTrainingLeaf
             .build(
-              reducedData.asInstanceOf[Seq[(Vector[AnyVal], Double, Double)]],
+              reducedData.asInstanceOf[Seq[TrainingRow[Double]]],
               GuessTheMeanLearner(),
               maxDepth - remainingDepth,
               rng
@@ -149,8 +149,8 @@ object MultiTaskTrainingNode {
         } else {
           ClassificationTrainingLeaf
             .build(
-              reducedData.asInstanceOf[Seq[(Vector[AnyVal], Char, Double)]],
-              GuessTheMeanLearner(),
+              reducedData.asInstanceOf[Seq[TrainingRow[Char]]],
+              GuessTheModeLearner(),
               maxDepth - remainingDepth,
               rng
             )
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 5eacb3e2..da841a82 100644
--- a/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTree.scala
+++ b/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTree.scala
@@ -1,12 +1,12 @@
 package io.citrine.lolo.trees.multitask
 
+import io.citrine.lolo.api.{Model, MultiTaskLearner, MultiTaskTrainingResult, ParallelModels, TrainingRow}
 import io.citrine.random.Random
 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.trees.splits.MultiTaskSplitter
-import io.citrine.lolo.{Model, MultiTaskLearner, MultiTaskTrainingResult, ParallelModels, PredictionResult}
 
 /**
   * A tree learner that operates on multiple labels.
@@ -27,16 +27,11 @@ case class MultiTaskTreeLearner(
     * Construct one regression or classification tree for each label.
     *
     * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
     * @param rng          random number generator for reproducibility
     * @return             sequence of models, one for each label
     */
-  override def train(
-      trainingData: Seq[(Vector[Any], Vector[Any])],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): MultiTaskTreeTrainingResult = {
-    val (inputs, labels) = trainingData.unzip
+  override def train(trainingData: Seq[TrainingRow[Vector[Any]]], rng: Random): MultiTaskTreeTrainingResult = {
+    val (inputs, labels, weights) = trainingData.map(_.asTuple).unzip3
     val repInput = inputs.head
     val repOutput = labels.head
     val labelIndices = repOutput.indices
@@ -65,16 +60,14 @@ case class MultiTaskTreeLearner(
 
     // Encode the inputs, outputs, and filter out zero weight rows
     val collectedData = inputs.indices
-      .map { i =>
-        (encodedInputs(i), encodedLabels(i), weights.map(_(i)).getOrElse(1.0))
-      }
-      .filter(_._3 > 0.0)
+      .map { i => TrainingRow(encodedInputs(i), encodedLabels(i), weights(i)) }
+      .filter(_.weight > 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
+      collectedData.head.inputs.size
     }
 
     // Construct the training tree
@@ -94,13 +87,13 @@ case class MultiTaskTreeLearner(
     // Stick the model trees into RegressionTree and ClassificationTree objects
     val models = labelIndices.map { i =>
       if (repOutput(i).isInstanceOf[Double]) {
-        new RegressionTree(
-          nodes(i).asInstanceOf[ModelNode[PredictionResult[Double]]],
+        RegressionTree(
+          nodes(i).asInstanceOf[ModelNode[Double]],
           inputEncoders
         )
       } else {
-        new ClassificationTree(
-          nodes(i).asInstanceOf[ModelNode[PredictionResult[Char]]],
+        ClassificationTree(
+          nodes(i).asInstanceOf[ModelNode[Char]],
           inputEncoders,
           outputEncoders(i).get
         )
@@ -114,27 +107,22 @@ case class MultiTaskTreeLearner(
       }
     }
 
-    new MultiTaskTreeTrainingResult(models, sumFeatureImportance)
+    MultiTaskTreeTrainingResult(models, sumFeatureImportance)
   }
 }
 
-class MultiTaskTreeTrainingResult(
-    models: Seq[Model[PredictionResult[Any]]],
-    featureImportance: Vector[Double]
+case class MultiTaskTreeTrainingResult(
+    override val models: Seq[Model[Any]],
+    nodeImportance: 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)
+  override val model: ParallelModels = ParallelModels(models, models.map(_.isInstanceOf[RegressionTree]))
 
+  override lazy val featureImportance: Option[Vector[Double]] = Some(
+    if (Math.abs(nodeImportance.sum) > 0) {
+      nodeImportance.map(_ / nodeImportance.sum)
+    } else {
+      nodeImportance.map(_ => 1.0 / nodeImportance.size)
+    }
+  )
 }
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 c013c46a..44c3ee7b 100644
--- a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingLeaf.scala
+++ b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingLeaf.scala
@@ -1,15 +1,15 @@
 package io.citrine.lolo.trees.regression
 
+import io.citrine.lolo.api.{Learner, TrainingResult, TrainingRow}
 import io.citrine.lolo.trees.TrainingLeaf
-import io.citrine.lolo.{Learner, TrainingResult}
 import io.citrine.random.Random
 
 import scala.collection.mutable
 
 /** Training leaf node for regression trees. */
 case class RegressionTrainingLeaf(
-    trainingData: Seq[(Vector[AnyVal], Double, Double)],
-    trainingResult: TrainingResult,
+    trainingData: Seq[TrainingRow[Double]],
+    trainingResult: TrainingResult[Double],
     depth: Int
 ) extends TrainingLeaf[Double] {
 
@@ -18,13 +18,13 @@ case class RegressionTrainingLeaf(
     *
     * @return feature importance as a vector
     */
-  override def featureImportance: scala.collection.mutable.ArraySeq[Double] = {
-    trainingResult.getFeatureImportance() match {
+  override def featureImportance: mutable.ArraySeq[Double] = {
+    trainingResult.featureImportance 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 (_, l, w) =>
+            case TrainingRow(_, 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))
@@ -32,7 +32,8 @@ case class RegressionTrainingLeaf(
         val impurity =
           Math.max(expectations._2 / expectations._3 - Math.pow(expectations._1 / expectations._3, 2.0), 0.0)
         mutable.ArraySeq(x: _*).map(_ * impurity)
-      case None => mutable.ArraySeq.fill(trainingData.head._1.size)(0.0)
+      case None =>
+        mutable.ArraySeq.fill(trainingData.head.inputs.size)(0.0)
     }
   }
 }
@@ -47,8 +48,8 @@ object RegressionTrainingLeaf {
     * @return a trained regression leaf node
     */
   def build(
-      trainingData: Seq[(Vector[AnyVal], Double, Double)],
-      leafLearner: Learner,
+      trainingData: Seq[TrainingRow[Double]],
+      leafLearner: Learner[Double],
       depth: Int,
       rng: Random
   ): RegressionTrainingLeaf =
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 d8dd98ca..3535828f 100644
--- a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingNode.scala
+++ b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingNode.scala
@@ -2,14 +2,14 @@ package io.citrine.lolo.trees.regression
 
 import io.citrine.lolo.trees.splits.{NoSplit, Split, Splitter}
 import io.citrine.lolo.trees.{InternalModelNode, ModelNode, TrainingNode}
-import io.citrine.lolo.{Learner, PredictionResult}
+import io.citrine.lolo.api.{Learner, TrainingRow}
 import io.citrine.random.Random
 
 case class RegressionTrainingNode(
-    trainingData: Seq[(Vector[AnyVal], Double, Double)],
+    trainingData: Seq[TrainingRow[Double]],
     leftNode: TrainingNode[Double],
     rightNode: TrainingNode[Double],
-    leafLearner: Learner,
+    leafLearner: Learner[Double],
     split: Split,
     deltaImpurity: Double
 ) extends TrainingNode[Double] {
@@ -19,13 +19,13 @@ case class RegressionTrainingNode(
     *
     * @return lightweight prediction node
     */
-  override def modelNode: ModelNode[PredictionResult[Double]] = {
-    new InternalModelNode[PredictionResult[Double]](
+  override def modelNode: ModelNode[Double] = {
+    InternalModelNode(
       split,
       leftNode.modelNode,
       rightNode.modelNode,
-      1,
-      trainingData.map(_._3).sum
+      outputDimension = 1,
+      trainingData.map(_.weight).sum
     )
   }
 
@@ -62,8 +62,8 @@ object RegressionTrainingNode {
     * @return the child node, either a RegressionTrainingNode or TrainingLeaf
     */
   def build(
-      trainingData: Seq[(Vector[AnyVal], Double, Double)],
-      leafLearner: Learner,
+      trainingData: Seq[TrainingRow[Double]],
+      leafLearner: Learner[Double],
       splitter: Splitter[Double],
       numFeatures: Int,
       minLeafInstances: Int,
@@ -73,7 +73,7 @@ object RegressionTrainingNode {
   ): TrainingNode[Double] = {
     val sufficientData = trainingData.size >= 2 * minLeafInstances &&
       remainingDepth > 0 &&
-      trainingData.exists(_._2 != trainingData.head._2)
+      trainingData.exists(_.label != trainingData.head.label)
     val (split: Split, deltaImpurity: Double) = if (sufficientData) {
       splitter.getBestSplit(trainingData, numFeatures, minLeafInstances, rng)
     } else {
@@ -83,7 +83,7 @@ object RegressionTrainingNode {
       case _: NoSplit =>
         RegressionTrainingLeaf.build(trainingData, leafLearner, maxDepth - remainingDepth, rng)
       case split: Split =>
-        val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r._1))
+        val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r.inputs))
         val leftNode = RegressionTrainingNode.build(
           trainingData = leftTrain,
           leafLearner = leafLearner,
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 bc8b6dcc..9ae3e219 100644
--- a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTree.scala
+++ b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTree.scala
@@ -1,12 +1,14 @@
 package io.citrine.lolo.trees.regression
 
 import breeze.linalg.DenseMatrix
+import io.citrine.lolo.api.{Learner, Model, PredictionResult, TrainingResult, TrainingRow}
 import io.citrine.random.Random
 import io.citrine.lolo.encoders.CategoricalEncoder
 import io.citrine.lolo.linear.GuessTheMeanLearner
 import io.citrine.lolo.trees.splits.{RegressionSplitter, Splitter}
 import io.citrine.lolo.trees.{ModelNode, TrainingNode, TreeMeta}
-import io.citrine.lolo.{Learner, Model, PredictionResult, TrainingResult}
+
+import scala.collection.mutable
 
 /**
   * Learner for regression trees
@@ -21,9 +23,9 @@ case class RegressionTreeLearner(
     numFeatures: Int = -1,
     maxDepth: Int = 30,
     minLeafInstances: Int = 1,
-    leafLearner: Option[Learner] = None,
+    leafLearner: Option[Learner[Double]] = None,
     splitter: Splitter[Double] = RegressionSplitter()
-) extends Learner {
+) extends Learner[Double] {
 
   /** Learner to use for training the leaves */
   @transient private lazy val myLeafLearner = leafLearner.getOrElse(GuessTheMeanLearner())
@@ -32,45 +34,31 @@ case class RegressionTreeLearner(
     * Train the tree by recursively partitioning (splitting) the training data on a single feature
     *
     * @param trainingData to train on
-    * @param weights      for the training rows, if applicable
     * @param rng          random number generator for reproducibility
     * @return a RegressionTree
     */
-  override def train(
-      trainingData: Seq[(Vector[Any], Any)],
-      weights: Option[Seq[Double]],
-      rng: Random
-  ): RegressionTreeTrainingResult = {
+  override def train(trainingData: Seq[TrainingRow[Double]], rng: Random): RegressionTreeTrainingResult = {
     require(trainingData.nonEmpty, s"The input training data was empty")
-    if (!trainingData.head._2.isInstanceOf[Double]) {
-      throw new IllegalArgumentException(
-        s"Tried to train regression on non-double labels, e.g.: ${trainingData.head._2}"
-      )
-    }
-    val repInput = trainingData.head._1
 
     /* Create encoders for any categorical features */
+    val repInput = trainingData.head.inputs
     val encoders: Seq[Option[CategoricalEncoder[Any]]] = repInput.zipWithIndex.map {
       case (v, i) =>
         if (v.isInstanceOf[Double]) {
           None
         } else {
-          Some(CategoricalEncoder.buildEncoder(trainingData.map(_._1(i))))
+          Some(CategoricalEncoder.buildEncoder(trainingData.map(_.inputs(i))))
         }
     }
 
     // Encode the training data.
-    val encodedTraining = trainingData.map(p => (CategoricalEncoder.encodeInput(p._1, encoders), p._2))
+    val encodedTraining = trainingData.map { row =>
+      val encodedInputs = CategoricalEncoder.encodeInput(row.inputs, encoders)
+      row.withInputs(encodedInputs)
+    }
 
     // Add the weights to the (features, label) tuples and remove any with zero weight.
-    val finalTraining = encodedTraining
-      .zip(weights.getOrElse(Seq.fill(trainingData.size)(1.0)))
-      .map {
-        case ((f, l), w) =>
-          (f, l.asInstanceOf[Double], w)
-      }
-      .filter(_._3 > 0)
-      .toVector
+    val finalTraining = encodedTraining.filter(_.weight > 0.0).toVector
 
     require(
       finalTraining.size >= 4,
@@ -81,7 +69,7 @@ case class RegressionTreeLearner(
     val numFeaturesActual = if (numFeatures > 0) {
       numFeatures
     } else {
-      finalTraining.head._1.size
+      finalTraining.head.inputs.size
     }
 
     // Recursively build the tree via its nodes and wrap the top node in a RegressionTreeTrainingResult.
@@ -95,33 +83,27 @@ case class RegressionTreeLearner(
       maxDepth = maxDepth,
       rng = rng
     )
-    new RegressionTreeTrainingResult(rootTrainingNode, encoders)
+    RegressionTreeTrainingResult(rootTrainingNode, encoders)
   }
 
 }
 
-class RegressionTreeTrainingResult(
+case class RegressionTreeTrainingResult(
     rootTrainingNode: TrainingNode[Double],
     encoders: Seq[Option[CategoricalEncoder[Any]]]
-) extends TrainingResult {
-  lazy val model = new RegressionTree(rootTrainingNode.modelNode, encoders)
-  lazy val importance = rootTrainingNode.featureImportance
-  private lazy val importanceNormalized = {
-    if (Math.abs(importance.sum) > 0) {
-      importance.map(_ / importance.sum)
-    } else {
-      importance.map(_ => 1.0 / importance.size)
-    }
-  }
+) extends TrainingResult[Double] {
 
-  override def getModel(): RegressionTree = model
+  override lazy val model: RegressionTree = RegressionTree(rootTrainingNode.modelNode, encoders)
 
-  /**
-    * Return the pre-computed influences
-    *
-    * @return feature influences as an array of doubles
-    */
-  override def getFeatureImportance(): Option[Vector[Double]] = Some(importanceNormalized.toVector)
+  lazy val nodeImportance: mutable.ArraySeq[Double] = rootTrainingNode.featureImportance
+
+  override lazy val featureImportance: Option[Vector[Double]] = Some(
+    if (Math.abs(nodeImportance.sum) > 0) {
+      nodeImportance.map(_ / nodeImportance.sum).toVector
+    } else {
+      nodeImportance.map(_ => 1.0 / nodeImportance.size).toVector
+    }
+  )
 }
 
 /**
@@ -130,10 +112,10 @@ class RegressionTreeTrainingResult(
   * @param root     of the tree
   * @param encoders for categorical variables
   */
-class RegressionTree(
-    root: ModelNode[PredictionResult[Double]],
+case class RegressionTree(
+    root: ModelNode[Double],
     encoders: Seq[Option[CategoricalEncoder[Any]]]
-) extends Model[RegressionTreeResult] {
+) extends Model[Double] {
 
   /**
     * Apply the model by calling predict and wrapping the results
@@ -141,8 +123,8 @@ class RegressionTree(
     * @param inputs to apply the model to
     * @return a prediction result which includes only the expected outputs
     */
-  override def transform(inputs: Seq[Vector[Any]]): RegressionTreeResult = {
-    new RegressionTreeResult(
+  override def transform(inputs: Seq[Vector[Any]]): RegressionTreePrediction = {
+    RegressionTreePrediction(
       inputs.map(inp => root.transform(CategoricalEncoder.encodeInput(inp, encoders)))
     )
   }
@@ -167,28 +149,27 @@ class RegressionTree(
   *
   * @param predictions sequence of predictions
   */
-class RegressionTreeResult(predictions: Seq[(PredictionResult[Double], TreeMeta)]) extends PredictionResult[Double] {
+case class RegressionTreePrediction(predictions: Seq[(PredictionResult[Double], TreeMeta)])
+    extends PredictionResult[Double] {
 
   /**
     * Get the predictions
     *
     * @return expected value of each prediction
     */
-  override def getExpected(): Seq[Double] = predictions.map(_._1.getExpected().head)
+  override def expected: Seq[Double] = predictions.map(_._1.expected.head)
 
   /**
     * Get the gradient or sensitivity of each prediction
     *
     * @return a vector of doubles for each prediction
     */
-  override def getGradient(): Option[Seq[Vector[Double]]] = {
-    if (!predictions.head._1.getGradient().isDefined) {
+  override def gradient: Option[Seq[Vector[Double]]] = {
+    if (predictions.head._1.gradient.isEmpty) {
       return None
     }
-    Some(predictions.map(_._1.getGradient().get.head))
+    Some(predictions.map(_._1.gradient.get.head))
   }
 
-  def getDepth(): Seq[Int] = {
-    predictions.map(_._2.depth)
-  }
+  def depth: Seq[Int] = predictions.map(_._2.depth)
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/BoltzmannSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/BoltzmannSplitter.scala
index 81f62db2..66e83918 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/BoltzmannSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/BoltzmannSplitter.scala
@@ -1,5 +1,6 @@
 package io.citrine.lolo.trees.splits
 
+import io.citrine.lolo.api.TrainingRow
 import io.citrine.random.Random
 import io.citrine.lolo.trees.impurity.VarianceCalculator
 import io.citrine.lolo.trees.splits.BoltzmannSplitter.SplitterResult
@@ -47,13 +48,13 @@ case class BoltzmannSplitter(temperature: Double) extends Splitter[Double] {
     * @return a split object that optimally divides data
     */
   def getBestSplit(
-      data: Seq[(Vector[AnyVal], Double, Double)],
+      data: Seq[TrainingRow[Double]],
       numFeatures: Int,
       minInstances: Int,
       rng: Random
   ): (Split, Double) = {
     /* Pre-compute these for the variance calculation */
-    val calculator = VarianceCalculator.build(data.map(_._2), data.map(_._3))
+    val calculator = VarianceCalculator.build(data.map(_.label), data.map(_.weight))
     val initialVariance = calculator.getImpurity
 
     // Don't split if there is no impurity to reduce
@@ -65,11 +66,11 @@ case class BoltzmannSplitter(temperature: Double) extends Splitter[Double] {
     val rep = data.head
 
     /* Try every feature index */
-    val featureIndices: Seq[Int] = rep._1.indices
+    val featureIndices: Seq[Int] = rep.inputs.indices
 
     val possibleSplits: Seq[SplitterResult] = rng.shuffle(featureIndices).take(numFeatures).flatMap { index =>
-      /* Use different spliters for each type */
-      rep._1(index) match {
+      /* Use different splitters for each type */
+      rep.inputs(index) match {
         case _: Double => BoltzmannSplitter.getBestRealSplit(data, calculator, index, minInstances, beta, rng)
         case _: Char   => BoltzmannSplitter.getBestCategoricalSplit(data, calculator, index, minInstances, beta, rng)
         case _: Any    => throw new IllegalArgumentException("Trying to split unknown feature type")
@@ -131,7 +132,7 @@ object BoltzmannSplitter {
     * @return the best split of this feature, along with its score, base, and result variance
     */
   def getBestRealSplit(
-      data: Seq[(Vector[AnyVal], Double, Double)],
+      data: Seq[TrainingRow[Double]],
       calculator: VarianceCalculator,
       index: Int,
       minCount: Int,
@@ -139,7 +140,7 @@ object BoltzmannSplitter {
       rng: Random
   ): Option[SplitterResult] = {
     /* 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 thinData = data.map(dat => (dat.inputs(index).asInstanceOf[Double], dat.label, dat.weight)).sortBy(_._1)
 
     /* Move the data from the right to the left partition one value at a time */
     calculator.reset()
@@ -192,7 +193,7 @@ object BoltzmannSplitter {
     * @return the best split of this feature, along with its score, base, and result variance
     */
   def getBestCategoricalSplit(
-      data: Seq[(Vector[AnyVal], Double, Double)],
+      data: Seq[TrainingRow[Double]],
       calculator: VarianceCalculator,
       index: Int,
       minCount: Int,
@@ -200,7 +201,7 @@ object BoltzmannSplitter {
       rng: Random
   ): Option[SplitterResult] = {
     /* Extract the features at the index */
-    val thinData = data.map(dat => (dat._1(index).asInstanceOf[Char], dat._2, dat._3))
+    val thinData = data.map(dat => (dat.inputs(index).asInstanceOf[Char], dat.label, dat.weight))
     val totalWeight = thinData.map(_._3).sum
 
     /* Group the data by categorical feature and compute the weighted sum and sum of the weights for each */
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 f8e1624e..548afa88 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/ClassificationSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/ClassificationSplitter.scala
@@ -1,5 +1,6 @@
 package io.citrine.lolo.trees.splits
 
+import io.citrine.lolo.api.TrainingRow
 import io.citrine.random.Random
 import io.citrine.lolo.trees.impurity.GiniCalculator
 
@@ -20,7 +21,7 @@ case class ClassificationSplitter(randomizedPivotLocation: Boolean = false) exte
     * @return a split object that optimally divides data
     */
   def getBestSplit(
-      data: Seq[(Vector[AnyVal], Char, Double)],
+      data: Seq[TrainingRow[Char]],
       numFeatures: Int,
       minInstances: Int,
       rng: Random
@@ -29,19 +30,18 @@ case class ClassificationSplitter(randomizedPivotLocation: Boolean = false) exte
     var bestImpurity = Double.MaxValue
 
     /* Pre-compute these for the variance calculation */
-    val calculator = GiniCalculator.build(data.map(p => (p._2, p._3)))
+    val calculator = GiniCalculator.build(data.map(p => (p.label, p.weight)))
     val initialImpurity = calculator.getImpurity
 
     val rep = data.head
     /* Try every feature index */
-    val featureIndices: Seq[Int] = rep._1.indices
+    val featureIndices: Seq[Int] = rep.inputs.indices
     rng.shuffle(featureIndices).take(numFeatures).foreach { index =>
       /* Use different spliters for each type */
-      val (possibleSplit, possibleImpurity) = rep._1(index) match {
-        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")
+      val (possibleSplit, possibleImpurity) = rep.inputs(index) match {
+        case _: Double => Splitter.getBestRealSplit(data, calculator, index, minInstances, randomizedPivotLocation, rng)
+        case _: Char   => getBestCategoricalSplit(data, calculator, index, minInstances)
+        case _: Any    => throw new IllegalArgumentException("Trying to split unknown feature type")
       }
 
       /* Keep track of the best split */
@@ -59,12 +59,12 @@ case class ClassificationSplitter(randomizedPivotLocation: Boolean = false) exte
   }
 
   def getBestCategoricalSplit(
-      data: Seq[(Vector[AnyVal], Char, Double)],
+      data: Seq[TrainingRow[Char]],
       calculator: GiniCalculator,
       index: Int,
       minCount: Int
   ): (CategoricalSplit, Double) = {
-    val thinData = data.map(dat => (dat._1(index).asInstanceOf[Char], dat._2, dat._3))
+    val thinData = data.map(dat => (dat.inputs(index).asInstanceOf[Char], dat.label, dat.weight))
     val groupedData = thinData.groupBy(_._1).view.mapValues { g =>
       val dict = g.groupBy(_._2).view.mapValues(v => v.map(_._3).sum)
       val impurity = dict.values.map(Math.pow(_, 2)).sum / Math.pow(dict.values.sum, 2)
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomClassificationSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomClassificationSplitter.scala
index 37b1a3fe..7b6f98ac 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomClassificationSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomClassificationSplitter.scala
@@ -1,5 +1,6 @@
 package io.citrine.lolo.trees.splits
 
+import io.citrine.lolo.api.TrainingRow
 import io.citrine.random.Random
 import io.citrine.lolo.trees.impurity.{GiniCalculator, ImpurityCalculator}
 
@@ -21,13 +22,13 @@ case class ExtraRandomClassificationSplitter() extends Splitter[Char] {
     * @return a split object that optimally divides data
     */
   def getBestSplit(
-      data: Seq[(Vector[AnyVal], Char, Double)],
+      data: Seq[TrainingRow[Char]],
       numFeatures: Int,
       minInstances: Int,
       rng: Random
   ): (Split, Double) = {
 
-    val calculator = GiniCalculator.build(data.map { p => (p._2, p._3) })
+    val calculator = GiniCalculator.build(data.map { p => (p.label, p.weight) })
     val initialVariance = calculator.getImpurity
     var bestSplit: Split = NoSplit()
     var bestVariance = Double.MaxValue
@@ -35,10 +36,10 @@ case class ExtraRandomClassificationSplitter() extends Splitter[Char] {
     val rep = data.head
 
     /* Try every feature index */
-    val featureIndices: Seq[Int] = rep._1.indices
+    val featureIndices: Seq[Int] = rep.inputs.indices
     rng.shuffle(featureIndices.toVector).take(numFeatures).foreach { index =>
       /* Use different spliters for each type */
-      val (possibleSplit, possibleVariance) = rep._1(index) match {
+      val (possibleSplit, possibleVariance) = rep.inputs(index) match {
         case _: Double => getBestRealSplit(data, calculator.copy(), index, minInstances, rng)
         case _: Char   => getBestCategoricalSplit(data, calculator.copy(), index, minInstances, rng)
         case _: Any    => throw new IllegalArgumentException("Trying to split unknown feature type")
@@ -68,14 +69,14 @@ case class ExtraRandomClassificationSplitter() extends Splitter[Char] {
     * @return the best split of this feature
     */
   def getBestRealSplit(
-      data: Seq[(Vector[AnyVal], Char, Double)],
+      data: Seq[TrainingRow[Char]],
       calculator: ImpurityCalculator[Char],
       index: Int,
       minCount: Int,
       rng: Random
   ): (RealSplit, Double) = {
     /* Pull out the feature that's considered here and sort by it */
-    val axis: Seq[Double] = data.map(_._1(index).asInstanceOf[Double])
+    val axis: Seq[Double] = data.map(_.inputs(index).asInstanceOf[Double])
     val lowerBound = axis.min
     val upperBound = axis.max
     val pivot = lowerBound + (upperBound - lowerBound) * rng.nextDouble()
@@ -85,8 +86,8 @@ case class ExtraRandomClassificationSplitter() extends Splitter[Char] {
     /* Move the data from the right to the left partition one value at a time */
     calculator.reset()
     data.foreach { dat =>
-      if (split.turnLeft(dat._1)) {
-        calculator.add(dat._2, dat._3)
+      if (split.turnLeft(dat.inputs)) {
+        calculator.add(dat.label, dat.weight)
       }
     }
 
@@ -103,14 +104,14 @@ case class ExtraRandomClassificationSplitter() extends Splitter[Char] {
     * @return the best split of this feature
     */
   def getBestCategoricalSplit(
-      data: Seq[(Vector[AnyVal], Char, Double)],
+      data: Seq[TrainingRow[Char]],
       calculator: ImpurityCalculator[Char],
       index: Int,
       minCount: Int,
       rng: Random
   ): (Split, Double) = {
     /* Extract the features at the index */
-    val thinData = data.map(dat => (dat._1(index).asInstanceOf[Char], dat._2, dat._3))
+    val thinData = data.map(dat => (dat.inputs(index).asInstanceOf[Char], dat.label, dat.weight))
 
     // Group the data by categorical feature and compute the weighted sum and sum of the weights for each
     val groupedData = thinData.groupBy(_._1).view.mapValues(g => (g.map(v => v._2 * v._3).sum, g.map(_._3).sum, g.size))
@@ -132,8 +133,8 @@ case class ExtraRandomClassificationSplitter() extends Splitter[Char] {
     /* Move the data from the right to the left partition one value at a time */
     calculator.reset()
     data.foreach { dat =>
-      if (split.turnLeft(dat._1)) {
-        calculator.add(dat._2, dat._3)
+      if (split.turnLeft(dat.inputs)) {
+        calculator.add(dat.label, dat.weight)
       }
     }
 
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitter.scala
index a813594a..8bb78921 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitter.scala
@@ -1,5 +1,6 @@
 package io.citrine.lolo.trees.splits
 
+import io.citrine.lolo.api.TrainingRow
 import io.citrine.random.Random
 import io.citrine.lolo.trees.impurity.{ImpurityCalculator, VarianceCalculator}
 
@@ -21,13 +22,13 @@ case class ExtraRandomRegressionSplitter() extends Splitter[Double] {
     * @return a split object that optimally divides data
     */
   def getBestSplit(
-      data: Seq[(Vector[AnyVal], Double, Double)],
+      data: Seq[TrainingRow[Double]],
       numFeatures: Int,
       minInstances: Int,
       rng: Random
   ): (Split, Double) = {
 
-    val calculator = VarianceCalculator.build(data.map(_._2), data.map(_._3))
+    val calculator = VarianceCalculator.build(data.map(_.label), data.map(_.weight))
     val initialVariance = calculator.getImpurity
     var bestSplit: Split = NoSplit()
     var bestVariance = Double.MaxValue
@@ -35,10 +36,10 @@ case class ExtraRandomRegressionSplitter() extends Splitter[Double] {
     val rep = data.head
 
     /* Try every feature index */
-    val featureIndices: Seq[Int] = rep._1.indices
+    val featureIndices: Seq[Int] = rep.inputs.indices
     rng.shuffle(featureIndices.toVector).take(numFeatures).foreach { index =>
       /* Use different spliters for each type */
-      val (possibleSplit, possibleVariance) = rep._1(index) match {
+      val (possibleSplit, possibleVariance) = rep.inputs(index) match {
         case _: Double => getBestRealSplit(data, calculator.copy(), index, minInstances, rng)
         case _: Char   => getBestCategoricalSplit(data, calculator.copy(), index, minInstances, rng)
         case _: Any    => throw new IllegalArgumentException("Trying to split unknown feature type")
@@ -68,14 +69,14 @@ case class ExtraRandomRegressionSplitter() extends Splitter[Double] {
     * @return the best split of this feature
     */
   def getBestRealSplit(
-      data: Seq[(Vector[AnyVal], Double, Double)],
+      data: Seq[TrainingRow[Double]],
       calculator: ImpurityCalculator[Double],
       index: Int,
       minCount: Int,
       rng: Random
   ): (RealSplit, Double) = {
     /* Pull out the feature that's considered here and sort by it */
-    val axis: Seq[Double] = data.map(_._1(index).asInstanceOf[Double])
+    val axis: Seq[Double] = data.map(_.inputs(index).asInstanceOf[Double])
     val lowerBound = axis.min
     val upperBound = axis.max
     val pivot = lowerBound + (upperBound - lowerBound) * rng.nextDouble()
@@ -85,8 +86,8 @@ case class ExtraRandomRegressionSplitter() extends Splitter[Double] {
     /* Move the data from the right to the left partition one value at a time */
     calculator.reset()
     data.foreach { dat =>
-      if (split.turnLeft(dat._1)) {
-        calculator.add(dat._2, dat._3)
+      if (split.turnLeft(dat.inputs)) {
+        calculator.add(dat.label, dat.weight)
       }
     }
 
@@ -103,14 +104,14 @@ case class ExtraRandomRegressionSplitter() extends Splitter[Double] {
     * @return the best split of this feature
     */
   def getBestCategoricalSplit(
-      data: Seq[(Vector[AnyVal], Double, Double)],
+      data: Seq[TrainingRow[Double]],
       calculator: ImpurityCalculator[Double],
       index: Int,
       minCount: Int,
       rng: Random
   ): (Split, Double) = {
     /* Extract the features at the index */
-    val thinData = data.map(dat => (dat._1(index).asInstanceOf[Char], dat._2, dat._3))
+    val thinData = data.map(dat => (dat.inputs(index).asInstanceOf[Char], dat.label, dat.weight))
 
     /* Group the data by categorical feature and compute the weighted sum and sum of the weights for each */
     val groupedData = thinData.groupBy(_._1).view.mapValues(g => (g.map(v => v._2 * v._3).sum, g.map(_._3).sum, g.size))
@@ -132,8 +133,8 @@ case class ExtraRandomRegressionSplitter() extends Splitter[Double] {
     /* Move the data from the right to the left partition one value at a time */
     calculator.reset()
     data.foreach { dat =>
-      if (split.turnLeft(dat._1)) {
-        calculator.add(dat._2, dat._3)
+      if (split.turnLeft(dat.inputs)) {
+        calculator.add(dat.label, dat.weight)
       }
     }
 
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 000c20af..b98ae919 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/MultiTaskSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/MultiTaskSplitter.scala
@@ -1,9 +1,10 @@
 package io.citrine.lolo.trees.splits
 
+import io.citrine.lolo.api.TrainingRow
 import io.citrine.random.Random
 import io.citrine.lolo.trees.impurity.MultiImpurityCalculator
 
-case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Splitter[Seq[AnyVal]] {
+case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Splitter[Vector[AnyVal]] {
 
   /**
     * Get the best split, considering numFeature random features (w/o replacement)
@@ -15,28 +16,27 @@ case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Sp
     * @return a split object that optimally divides data
     */
   def getBestSplit(
-      data: Seq[(Vector[AnyVal], Seq[AnyVal], Double)],
+      data: Seq[TrainingRow[Vector[AnyVal]]],
       numFeatures: Int,
       minInstances: Int,
       rng: Random
   ): (Split, Double) = {
     var bestSplit: Split = NoSplit()
     var bestImpurity = Double.MaxValue
-    val calculator = MultiImpurityCalculator.build(data.map(_._2), data.map(_._3))
+    val calculator = MultiImpurityCalculator.build(data.map(_.label), data.map(_.weight))
     val initialImpurity = calculator.getImpurity
 
     // Pre-compute these for the variance calculation.
     val rep = data.head
 
     // Try every feature index.
-    val featureIndices: Seq[Int] = rep._1.indices
+    val featureIndices: Seq[Int] = rep.inputs.indices
     rng.shuffle(featureIndices).take(numFeatures).foreach { index =>
       // Use different splitters for each type.
-      val (possibleSplit, possibleImpurity) = rep._1(index) match {
-        case _: Double =>
-          Splitter.getBestRealSplit[Seq[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")
+      val (possibleSplit, possibleImpurity) = rep.inputs(index) match {
+        case _: Double => Splitter.getBestRealSplit(data, calculator, index, minInstances, randomizePivotLocation, rng)
+        case _: Char   => getBestCategoricalSplit(data, calculator, index, minInstances)
+        case _: Any    => throw new IllegalArgumentException("Trying to split unknown feature type")
       }
 
       // Keep track of the best split.
@@ -63,14 +63,13 @@ case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Sp
     * @return the best split of this feature
     */
   def getBestCategoricalSplit(
-      data: Seq[(Vector[AnyVal], Seq[AnyVal], Double)],
+      data: Seq[TrainingRow[Vector[AnyVal]]],
       calculator: MultiImpurityCalculator,
       index: Int,
       minInstances: Int
   ): (Split, Double) = {
     // Extract the features at the index.
-    val thinData: Seq[(Char, Seq[AnyVal], Double)] =
-      data.map(dat => (dat._1(index).asInstanceOf[Char], dat._2, dat._3))
+    val thinData = data.map(dat => (dat.inputs(index).asInstanceOf[Char], dat.label, dat.weight))
     val totalWeight = thinData.map(_._3).sum
 
     // Group the data by categorical feature and compute the weighted sum and sum of the weights for each.
@@ -120,7 +119,7 @@ case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Sp
     * @param labels is a seq of (Array of multiple labels, single weight)
     * @return the impurity, which is in [0, number of labels * sum of weights]
     */
-  def computeImpurity(labels: Seq[(Seq[AnyVal], Double)]): Double = {
+  def computeImpurity(labels: Seq[(Vector[AnyVal], Double)]): Double = {
     val calculator = MultiImpurityCalculator.build(labels.map(_._1), labels.map(_._2))
     calculator.getImpurity
   }
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 da98baf4..c2ba821d 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/RegressionSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/RegressionSplitter.scala
@@ -1,5 +1,6 @@
 package io.citrine.lolo.trees.splits
 
+import io.citrine.lolo.api.TrainingRow
 import io.citrine.lolo.trees.impurity.VarianceCalculator
 import io.citrine.random.Random
 
@@ -32,13 +33,13 @@ case class RegressionSplitter(randomizePivotLocation: Boolean = false) extends S
     * @return a split object that optimally divides data
     */
   def getBestSplit(
-      data: Seq[(Vector[AnyVal], Double, Double)],
+      data: Seq[TrainingRow[Double]],
       numFeatures: Int,
       minInstances: Int,
       rng: Random
   ): (Split, Double) = {
 
-    val calculator = VarianceCalculator.build(data.map(_._2), data.map(_._3))
+    val calculator = VarianceCalculator.build(data.map(_.label), data.map(_.weight))
     val initialVariance = calculator.getImpurity
     var bestSplit: Split = NoSplit()
     var bestVariance = Double.MaxValue
@@ -46,10 +47,10 @@ case class RegressionSplitter(randomizePivotLocation: Boolean = false) extends S
     val rep = data.head
 
     /* Try every feature index */
-    val featureIndices: Seq[Int] = rep._1.indices
+    val featureIndices: Seq[Int] = rep.inputs.indices
     rng.shuffle(featureIndices).take(numFeatures).foreach { index =>
       /* Use different splitters for each type */
-      val (possibleSplit, possibleVariance) = rep._1(index) match {
+      val (possibleSplit, possibleVariance) = rep.inputs(index) match {
         case _: Double =>
           Splitter.getBestRealSplit[Double](data, calculator.copy(), index, minInstances, randomizePivotLocation, rng)
         case _: Char => getBestCategoricalSplit(data, calculator.copy(), index, minInstances)
@@ -78,13 +79,13 @@ case class RegressionSplitter(randomizePivotLocation: Boolean = false) extends S
     * @return the best split of this feature
     */
   def getBestCategoricalSplit(
-      data: Seq[(Vector[AnyVal], Double, Double)],
+      data: Seq[TrainingRow[Double]],
       calculator: VarianceCalculator,
       index: Int,
       minCount: Int
   ): (CategoricalSplit, Double) = {
     /* Extract the features at the index */
-    val thinData = data.map(dat => (dat._1(index).asInstanceOf[Char], dat._2, dat._3))
+    val thinData = data.map(dat => (dat.inputs(index).asInstanceOf[Char], dat.label, dat.weight))
     val totalWeight = thinData.map(_._3).sum
 
     /* Group the data by categorical feature and compute the weighted sum and sum of the weights for each */
@@ -99,7 +100,7 @@ case class RegressionSplitter(randomizePivotLocation: Boolean = false) extends S
     /* Compute the average label for each categorical value */
     val categoryAverages: Map[Char, Double] = groupedData.mapValues(p => p._1 / p._2).toMap
 
-    /* Create an orderd list of the categories by average label */
+    /* Create an ordered list of the categories by average label */
     val orderedNames = categoryAverages.toSeq.sortBy(_._2).map(_._1)
 
     /* Base cases for the iteration */
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/Split.scala b/src/main/scala/io/citrine/lolo/trees/splits/Split.scala
index 0c2a5b2b..f61f145e 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/Split.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/Split.scala
@@ -13,7 +13,7 @@ sealed trait Split extends Serializable {
     * @param input vector of any type
     * @return true if input takes the left split
     */
-  def turnLeft(input: Vector[AnyVal]): Boolean
+  def turnLeft(input: Vector[Any]): Boolean
 
   /**
     * Get the index of the input vector that is used to pick this split
@@ -34,7 +34,7 @@ case class NoSplit() extends Split {
     * @param input vector of any type
     * @return true if input takes the left split
     */
-  override def turnLeft(input: Vector[AnyVal]): Boolean = false
+  override def turnLeft(input: Vector[Any]): Boolean = false
 
   /**
     * Get the index of the input vector that is used to pick this split
@@ -58,7 +58,7 @@ case class RealSplit(index: Int, pivot: Double) extends Split {
     * @param input vector of any type
     * @return true if input takes the left split
     */
-  override def turnLeft(input: Vector[AnyVal]): Boolean = {
+  override def turnLeft(input: Vector[Any]): Boolean = {
     if (pivot.isNaN) {
       // the normal sort order for scala puts NaN >> any other double
       // so if the pivot is NaN and the value is not, then it is 'less than' the pivot
@@ -90,7 +90,7 @@ case class CategoricalSplit(index: Int, includeSet: BitSet) extends Split {
     * @param input vector of any type
     * @return true if input takes the left split
     */
-  override def turnLeft(input: Vector[AnyVal]): Boolean = {
+  override def turnLeft(input: Vector[Any]): Boolean = {
     includeSet.contains(input(index).asInstanceOf[Char].toInt)
   }
 }
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 3b07d693..b82c521f 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/Splitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/Splitter.scala
@@ -1,11 +1,12 @@
 package io.citrine.lolo.trees.splits
 
+import io.citrine.lolo.api.TrainingRow
 import io.citrine.lolo.trees.impurity.ImpurityCalculator
 import io.citrine.random.Random
 
 trait Splitter[T] {
   def getBestSplit(
-      data: Seq[(Vector[AnyVal], T, Double)],
+      data: Seq[TrainingRow[T]],
       numFeatures: Int,
       minInstances: Int,
       rng: Random = Random()
@@ -37,7 +38,7 @@ object Splitter {
     * @return the best split of this feature
     */
   def getBestRealSplit[T](
-      data: Seq[(Vector[AnyVal], T, Double)],
+      data: Seq[TrainingRow[T]],
       calculator: ImpurityCalculator[T],
       index: Int,
       minCount: Int,
@@ -45,7 +46,7 @@ object Splitter {
       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 thinData = data.map(dat => (dat.inputs(index).asInstanceOf[Double], dat.label, dat.weight)).sortBy(_._1)
     val features = thinData.map(x => x._1)
 
     var bestImpurity = Double.MaxValue
diff --git a/src/main/scala/io/citrine/lolo/util/LoloPyDataLoader.scala b/src/main/scala/io/citrine/lolo/util/LoloPyDataLoader.scala
index cfc58fe3..571d2f66 100644
--- a/src/main/scala/io/citrine/lolo/util/LoloPyDataLoader.scala
+++ b/src/main/scala/io/citrine/lolo/util/LoloPyDataLoader.scala
@@ -4,7 +4,7 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream,
 import java.nio.ByteBuffer
 import java.nio.ByteOrder
 import java.util.zip._
-import io.citrine.lolo.{MultiTaskModelPredictionResult, PredictionResult}
+import io.citrine.lolo.api.{MultiTaskModelPredictionResult, PredictionResult}
 
 /**
   * Tool used to transfer data from LoloPy to the JVM
@@ -76,7 +76,7 @@ object LoloPyDataLoader {
     * @return Byte array of doubles in native system order
     */
   def getRegressionExpected(predictionResult: PredictionResult[Any]): Array[Byte] = {
-    val predResults: Seq[Double] = predictionResult.getExpected().asInstanceOf[Seq[Double]]
+    val predResults: Seq[Double] = predictionResult.expected.asInstanceOf[Seq[Double]]
     send1DArray(predResults)
   }
 
@@ -86,7 +86,7 @@ object LoloPyDataLoader {
     * @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
+    val predResults = predictionResult.expected.asInstanceOf[Seq[Seq[Double]]].flatten
     send1DArray(predResults)
   }
 
@@ -96,7 +96,7 @@ object LoloPyDataLoader {
     * @return Byte of array of doubles in native system order
     */
   def getImportanceScores(predictionResult: PredictionResult[Any]): Array[Byte] = {
-    send1DArray(predictionResult.getImportanceScores().get.flatten)
+    send1DArray(predictionResult.importanceScores.get.flatten)
   }
 
   /**
@@ -105,7 +105,7 @@ object LoloPyDataLoader {
     * @return Byte array of doubles in native system order
     */
   def getRegressionUncertainty(predictionResult: PredictionResult[Any]): Array[Byte] = {
-    val predResults: Seq[Double] = predictionResult.getUncertainty().get.asInstanceOf[Seq[Double]]
+    val predResults: Seq[Double] = predictionResult.uncertainty().get.asInstanceOf[Seq[Double]]
     send1DArray(predResults)
   }
 
@@ -115,7 +115,7 @@ object LoloPyDataLoader {
     * @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
+    val uncertaintyResults = predictionResult.uncertainty().get.asInstanceOf[Seq[Seq[Double]]].flatten
     send1DArray(uncertaintyResults)
   }
 
@@ -130,7 +130,7 @@ object LoloPyDataLoader {
     * @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
+    val correlationResults = predictionResult.uncertaintyCorrelation(i, j).get
     send1DArray(correlationResults)
   }
 
@@ -154,7 +154,7 @@ object LoloPyDataLoader {
     * @return Bytes of a integer array of the predicted class labels
     */
   def getClassifierExpected(predictionResult: PredictionResult[Any]): Array[Byte] = {
-    val expect = predictionResult.getExpected().asInstanceOf[Seq[Int]]
+    val expect = predictionResult.expected.asInstanceOf[Seq[Int]]
     val buffer = ByteBuffer.allocate(expect.length * 4).order(ByteOrder.nativeOrder())
     expect.foreach(buffer.putInt)
     buffer.array
@@ -173,7 +173,7 @@ object LoloPyDataLoader {
     // Get an iterator over the number of classes
     val classes = 0 until nClasses
     val probs = predictionResult
-      .getUncertainty()
+      .uncertainty()
       .get
       .asInstanceOf[Seq[Map[Int, Double]]]
       .map(x => classes.map(i => x.getOrElse(i, 0.0)))
@@ -217,14 +217,17 @@ object LoloPyDataLoader {
 
   /**
     * Create a PredictionResult object from the mean and uncertainty
-    * @param expected Mean of the predictions of a model
-    * @param uncertainty Uncertainty of the predictions
+    * @param thisExpected    Mean of the predictions of a model
+    * @param thisUncertainty Uncertainty of the predictions
     * @return Prediction result object
     */
-  def makeRegressionPredictionResult(expected: Seq[Double], uncertainty: Seq[Double]): PredictionResult[Double] = {
+  def makeRegressionPredictionResult(
+      thisExpected: Seq[Double],
+      thisUncertainty: Seq[Double]
+  ): PredictionResult[Double] = {
     new PredictionResult[Double] {
-      override def getExpected(): Seq[Double] = expected
-      override def getUncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = Some(uncertainty)
+      override def expected: Seq[Double] = thisExpected
+      override def uncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = Some(thisUncertainty)
     }
   }
 }
diff --git a/src/main/scala/io/citrine/lolo/validation/CrossValidation.scala b/src/main/scala/io/citrine/lolo/validation/CrossValidation.scala
index d1e2b82e..9ef6f9e9 100644
--- a/src/main/scala/io/citrine/lolo/validation/CrossValidation.scala
+++ b/src/main/scala/io/citrine/lolo/validation/CrossValidation.scala
@@ -1,7 +1,7 @@
 package io.citrine.lolo.validation
 
+import io.citrine.lolo.api.{Learner, PredictionResult, TrainingRow}
 import io.citrine.random.Random
-import io.citrine.lolo.{Learner, PredictionResult}
 
 /**
   * Methods tha use cross-validation to calculate predicted-vs-actual data and metric estimates
@@ -21,8 +21,8 @@ case object CrossValidation {
     * @return a Map from the name of the metric to its mean value and the error in that mean
     */
   def kFoldCrossvalidation[T](
-      trainingData: Seq[(Vector[Any], T)],
-      learner: Learner,
+      trainingData: Seq[TrainingRow[T]],
+      learner: Learner[T],
       metrics: Map[String, Merit[T]],
       k: Int = 8,
       nTrial: Int = 1,
@@ -46,25 +46,24 @@ case object CrossValidation {
     * @return an iterable over predicted-vs-actual for each fold
     */
   def kFoldPvA[T](
-      trainingData: Seq[(Vector[Any], T)],
-      learner: Learner,
+      trainingData: Seq[TrainingRow[T]],
+      learner: Learner[T],
       k: Int = 8,
       nTrial: Int = 1,
       rng: Random = Random()
   ): Iterable[(PredictionResult[T], Seq[T])] = {
     val nTest: Int = Math.ceil(trainingData.size.toDouble / k).toInt
     (0 until nTrial).flatMap { _ =>
-      val folds: Seq[Seq[(Vector[Any], T)]] = rng.shuffle(trainingData).grouped(nTest).toSeq
+      val folds = rng.shuffle(trainingData).grouped(nTest).toSeq
 
       folds.indices.map { idx =>
         val (testFolds, trainFolds) = folds.zipWithIndex.partition(_._2 == idx)
         val testData = testFolds.flatMap(_._1)
         val trainData = trainFolds.flatMap(_._1)
-        val model = learner.train(trainData, rng = rng).getModel()
-        val predictions: PredictionResult[T] = model.transform(testData.map(_._1)).asInstanceOf[PredictionResult[T]]
-        (predictions, testData.map(_._2))
+        val model = learner.train(trainData, rng = rng).model
+        val predictions: PredictionResult[T] = model.transform(testData.map(_.inputs))
+        (predictions, testData.map(_.label))
       }
     }
   }
-
 }
diff --git a/src/main/scala/io/citrine/lolo/validation/Merit.scala b/src/main/scala/io/citrine/lolo/validation/Merit.scala
index 99018448..ec8e6d1c 100644
--- a/src/main/scala/io/citrine/lolo/validation/Merit.scala
+++ b/src/main/scala/io/citrine/lolo/validation/Merit.scala
@@ -1,7 +1,6 @@
 package io.citrine.lolo.validation
 
-import java.util
-import io.citrine.lolo.PredictionResult
+import io.citrine.lolo.api.PredictionResult
 import io.citrine.random.Random
 import org.knowm.xchart.XYChart
 
@@ -44,13 +43,12 @@ case object RootMeanSquareError extends Merit[Double] {
       rng: Random = Random()
   ): Double = {
     Math.sqrt(
-      predictionResult
-        .getExpected()
+      predictionResult.expected
         .zip(actual)
         .map {
           case (x, y) => Math.pow(x - y, 2)
         }
-        .sum / predictionResult.getExpected().size
+        .sum / predictionResult.expected.size
     )
   }
 }
@@ -66,7 +64,7 @@ case object CoefficientOfDetermination extends Merit[Double] {
   ): Double = {
     val averageActual = actual.sum / actual.size
     val sumOfSquares = actual.map(x => Math.pow(x - averageActual, 2)).sum
-    val sumOfResiduals = predictionResult.getExpected().zip(actual).map { case (x, y) => Math.pow(x - y, 2.0) }.sum
+    val sumOfResiduals = predictionResult.expected.zip(actual).map { case (x, y) => Math.pow(x - y, 2.0) }.sum
     1.0 - sumOfResiduals / sumOfSquares
   }
 }
@@ -80,11 +78,11 @@ case object StandardConfidence extends Merit[Double] {
       actual: Seq[Double],
       rng: Random = Random()
   ): Double = {
-    if (predictionResult.getUncertainty().isEmpty) return 0.0
+    if (predictionResult.uncertainty().isEmpty) return 0.0
 
-    predictionResult.getExpected().lazyZip(predictionResult.getUncertainty().get).lazyZip(actual).count {
+    predictionResult.expected.lazyZip(predictionResult.uncertainty().get).lazyZip(actual).count {
       case (x, sigma: Double, y) => Math.abs(x - y) < sigma
-    } / predictionResult.getExpected().size.toDouble
+    } / predictionResult.expected.size.toDouble
   }
 }
 
@@ -97,9 +95,9 @@ case class StandardError(rescale: Double = 1.0) extends Merit[Double] {
       actual: Seq[Double],
       rng: Random = Random()
   ): Double = {
-    if (predictionResult.getUncertainty().isEmpty) return Double.PositiveInfinity
+    if (predictionResult.uncertainty().isEmpty) return Double.PositiveInfinity
     val standardized =
-      predictionResult.getExpected().lazyZip(predictionResult.getUncertainty().get).lazyZip(actual).map {
+      predictionResult.expected.lazyZip(predictionResult.uncertainty().get).lazyZip(actual).map {
         case (x, sigma: Double, y) => (x - y) / sigma
       }
     rescale * Math.sqrt(standardized.map(Math.pow(_, 2.0)).sum / standardized.size)
@@ -123,12 +121,9 @@ case object UncertaintyCorrelation extends Merit[Double] {
       actual: Seq[Double],
       rng: Random = Random()
   ): Double = {
-    val predictedUncertaintyActual: Seq[(Double, Double, Double)] = (
-      predictionResult
-        .getExpected()
-        .lazyZip(predictionResult.getUncertainty().get.asInstanceOf[Seq[Double]])
-        .lazyZip(actual)
-      )
+    val predictedUncertaintyActual: Seq[(Double, Double, Double)] = predictionResult.expected
+      .lazyZip(predictionResult.uncertainty().get.asInstanceOf[Seq[Double]])
+      .lazyZip(actual)
       .toSeq
 
     val ideal = predictedUncertaintyActual.map {
diff --git a/src/main/scala/io/citrine/lolo/validation/StatisticalValidation.scala b/src/main/scala/io/citrine/lolo/validation/StatisticalValidation.scala
index 4a39a076..49198072 100644
--- a/src/main/scala/io/citrine/lolo/validation/StatisticalValidation.scala
+++ b/src/main/scala/io/citrine/lolo/validation/StatisticalValidation.scala
@@ -1,6 +1,6 @@
 package io.citrine.lolo.validation
 
-import io.citrine.lolo.{Learner, PredictionResult}
+import io.citrine.lolo.api.{Learner, PredictionResult, TrainingRow}
 import io.citrine.random.Random
 
 /**
@@ -28,19 +28,19 @@ case class StatisticalValidation() {
     * @return predicted-vs-actual data that can be fed into a metric or visualization
     */
   def generativeValidation[T](
-      source: Iterator[(Vector[Any], T)],
-      learner: Learner,
+      source: Iterator[TrainingRow[T]],
+      learner: Learner[T],
       nTrain: Int,
       nTest: Int,
       nRound: Int,
       rng: Random
   ): Iterator[(PredictionResult[T], Seq[T])] = {
     Iterator.tabulate(nRound) { _ =>
-      val trainingData: Seq[(Vector[Any], T)] = source.take(nTrain).toSeq
-      val model = learner.train(trainingData, rng = rng).getModel()
-      val testData: Seq[(Vector[Any], T)] = source.take(nTest).toSeq
-      val predictions: PredictionResult[T] = model.transform(testData.map(_._1)).asInstanceOf[PredictionResult[T]]
-      (predictions, testData.map(_._2))
+      val trainingData = source.take(nTrain).toSeq
+      val model = learner.train(trainingData, rng = rng).model
+      val testData = source.take(nTest).toSeq
+      val predictions: PredictionResult[T] = model.transform(testData.map(_.inputs))
+      (predictions, testData.map(_.label))
     }
   }
 
@@ -64,8 +64,8 @@ case class StatisticalValidation() {
     * @return predicted-vs-actual data that can be fed into a metric or visualization
     */
   def generativeValidation[T](
-      source: Iterable[(Vector[Any], T)],
-      learner: Learner,
+      source: Iterable[TrainingRow[T]],
+      learner: Learner[T],
       nTrain: Int,
       nTest: Int,
       nRound: Int,
@@ -73,11 +73,10 @@ case class StatisticalValidation() {
   ): Iterator[(PredictionResult[T], Seq[T])] = {
     Iterator.tabulate(nRound) { _ =>
       val subset = rng.shuffle(source).take(nTrain + nTest)
-      val (trainingData: Seq[(Vector[Any], T)], testData: Seq[(Vector[Any], T)]) = subset.splitAt(nTrain)
-      val model = learner.train(trainingData, rng = rng).getModel()
-      val predictions: PredictionResult[T] = model.transform(testData.map(_._1)).asInstanceOf[PredictionResult[T]]
-      (predictions, testData.map(_._2))
+      val (trainingData, testData) = subset.toVector.splitAt(nTrain)
+      val model = learner.train(trainingData, rng = rng).model
+      val predictions: PredictionResult[T] = model.transform(testData.map(_.inputs))
+      (predictions, testData.map(_.label))
     }
   }
-
 }
diff --git a/src/main/scala/io/citrine/lolo/validation/Visualization.scala b/src/main/scala/io/citrine/lolo/validation/Visualization.scala
index cc878430..707ea878 100644
--- a/src/main/scala/io/citrine/lolo/validation/Visualization.scala
+++ b/src/main/scala/io/citrine/lolo/validation/Visualization.scala
@@ -1,6 +1,6 @@
 package io.citrine.lolo.validation
 
-import io.citrine.lolo.PredictionResult
+import io.citrine.lolo.api.PredictionResult
 import org.apache.commons.math3.distribution.CauchyDistribution
 import org.knowm.xchart.XYSeries.XYSeriesRenderStyle
 import org.knowm.xchart._
@@ -39,9 +39,8 @@ case class StandardResidualHistogram(
   override def visualize(data: Iterable[(PredictionResult[Double], Seq[Double])]): CategoryChart = {
     val pua: Seq[(Double, Double, Double)] = data.flatMap {
       case (predictions, actual) =>
-        predictions
-          .getExpected()
-          .lazyZip(predictions.getUncertainty().get.asInstanceOf[Seq[Double]])
+        predictions.expected
+          .lazyZip(predictions.uncertainty().get.asInstanceOf[Seq[Double]])
           .lazyZip(actual)
           .toSeq
     }.toSeq
@@ -90,7 +89,7 @@ case class PredictedVsActual() extends Visualization[Double] {
 
     val flattened: Iterable[(Double, Double, Double)] = data.flatMap {
       case (pred, actual: Seq[Double]) =>
-        actual.lazyZip(pred.getExpected()).lazyZip(pred.getUncertainty().get.asInstanceOf[Seq[Double]]).toSeq
+        actual.lazyZip(pred.expected).lazyZip(pred.uncertainty().get.asInstanceOf[Seq[Double]]).toSeq
     }
 
     val actual = flattened.map(_._1).toArray
@@ -123,8 +122,8 @@ case class ErrorVsUncertainty(magnitude: Boolean = true) extends Visualization[D
 
     val flattened: Iterable[(Double, Double)] = data.flatMap {
       case (pred, actual: Seq[Double]) =>
-        val sigmas = pred.getUncertainty().get.asInstanceOf[Seq[Double]]
-        val errors = actual.zip(pred.getExpected()).map {
+        val sigmas = pred.uncertainty().get.asInstanceOf[Seq[Double]]
+        val errors = actual.zip(pred.expected).map {
           case (x, y) if magnitude => Math.abs(x - y)
           case (x, y)              => y - x
         }


Test patch:
diff --git a/python/lolopy/tests/test_learners.py b/python/lolopy/tests/test_learners.py
index 696db9e5..4f9f5281 100644
--- a/python/lolopy/tests/test_learners.py
+++ b/python/lolopy/tests/test_learners.py
@@ -1,6 +1,7 @@
 from lolopy.learners import (
     RandomForestRegressor,
     RandomForestClassifier,
+    MultiTaskRandomForest,
     RegressionTreeLearner,
     LinearRegression,
     ExtraRandomTreesRegressor,
@@ -13,6 +14,9 @@
 import pickle as pkl
 import numpy as np
 
+import logging
+logging.getLogger("py4j.java_gateway").setLevel(logging.ERROR)
+
 
 def _make_linear_data():
     """Make data corresponding to y = x + 1
@@ -91,7 +95,7 @@ def test_reproducibility(self):
         self.assertTrue((pred1 == pred2).all())
 
     def test_rf_multioutput_regressor(self):
-        rf = RandomForestRegressor()
+        rf = MultiTaskRandomForest()
         # A regression dataset with 3 outputs
         X, y = load_linnerud(return_X_y=True)
         num_data = len(X)
@@ -231,7 +235,8 @@ def test_adjust_rtree_learners(self):
 class TestExtraRandomTrees(TestCase):
 
     def test_extra_random_trees_regressor(self):
-        rf = ExtraRandomTreesRegressor()
+        # Setting disable_bootstrap=False allows us to generate uncertainty estimates from the bagged ensemble
+        rf = ExtraRandomTreesRegressor(disable_bootstrap=False)
 
         # Train the model
         X, y = load_diabetes(return_X_y=True)
@@ -264,6 +269,7 @@ def test_extra_random_trees_regressor(self):
         self.assertAlmostEqual(1.0, np.sum(rf.feature_importances_))
 
         # Run predictions with std dev
+        # Requires disable_bootstrap=False on the learner to generate std dev alongside predictions
         y_pred, y_std = rf.predict(X, return_std=True)
         self.assertEqual(len(y_pred), len(y_std))
         self.assertTrue((y_std >= 0).all())  # They must be positive
diff --git a/src/test/scala/io/citrine/lolo/AccuracyTest.scala b/src/test/scala/io/citrine/lolo/AccuracyTest.scala
index 84135a14..1d4a31bc 100644
--- a/src/test/scala/io/citrine/lolo/AccuracyTest.scala
+++ b/src/test/scala/io/citrine/lolo/AccuracyTest.scala
@@ -1,6 +1,7 @@
 package io.citrine.lolo
 
-import io.citrine.lolo.bags.Bagger
+import io.citrine.lolo.api.{Learner, TrainingRow}
+import io.citrine.lolo.bags.{Bagger, RegressionBagger}
 import io.citrine.lolo.trees.regression.RegressionTreeLearner
 import io.citrine.lolo.trees.splits.{BoltzmannSplitter, RegressionSplitter}
 import io.citrine.random.Random
@@ -17,16 +18,12 @@ class AccuracyTest extends SeedRandomMixIn {
   val nScal: Int = 2
   val minInstances: Int = 1
 
-  val trainingData: Seq[(Vector[Any], Double)] = TestUtils
-    .binTrainingData(
-      TestUtils.generateTrainingData(nRow, nFeat, noise = noiseLevel, rng = rng),
-      inputBins = Seq((4, 32))
-    )
-    .asInstanceOf[Seq[(Vector[Any], Double)]] // binTrainingData isn't binning the labels
+  val trainingData: Seq[TrainingRow[Double]] =
+    DataGenerator.generate(nRow, nFeat, noise = noiseLevel, rng = rng).withBinnedInputs(bins = Seq((4, 32))).data
 
   // Get the out-of-bag RMSE
-  private def computeMetrics(learner: Learner, rng: Random): Double = {
-    learner.train(trainingData, rng = rng).getLoss().get
+  private def computeMetrics(learner: Learner[Double], rng: Random): Double = {
+    learner.train(trainingData, rng = rng).loss.get
   }
 
   /**
@@ -35,7 +32,7 @@ class AccuracyTest extends SeedRandomMixIn {
   @Test
   def testRandomForest(): Unit = {
     val baseLearner = RegressionTreeLearner(numFeatures = nFeat / 3, minLeafInstances = minInstances)
-    val learner = new Bagger(baseLearner, numBags = nRow * nScal)
+    val learner = RegressionBagger(baseLearner, numBags = nRow * nScal)
     val error = computeMetrics(learner, rng)
     assert(error > noiseLevel, s"Can't do better than noise")
     assert(error < 4.0, "Error increased, probably due to a change in configuration")
@@ -51,8 +48,7 @@ class AccuracyTest extends SeedRandomMixIn {
         numFeatures = nFeat,
         splitter = RegressionSplitter(randomizePivotLocation = true)
       )
-      val learner =
-        new Bagger(baseLearner, numBags = nRow * 16)
+      val learner = RegressionBagger(baseLearner, numBags = nRow * 16)
       // println(s"Normal train time: ${Stopwatch.time(computeMetrics(learner))}")
       computeMetrics(learner, rng)
     }
@@ -61,8 +57,7 @@ class AccuracyTest extends SeedRandomMixIn {
         numFeatures = nFeat,
         splitter = BoltzmannSplitter(temperature = Float.MinPositiveValue)
       )
-      val learner =
-        new Bagger(baseLearner, numBags = nRow * 16)
+      val learner = RegressionBagger(baseLearner, numBags = nRow * 16)
       // println(s"Annealing train time: ${Stopwatch.time(computeMetrics(learner))}")
       computeMetrics(learner, rng)
     }
@@ -78,16 +73,11 @@ class AccuracyTest extends SeedRandomMixIn {
   * Driver code to study the performance vs temperature
   *
   * This isn't cast as a test, but can be used to try to understand the behavior of Boltzmann trees on some simple problems.
-  * TODO: turn this into a demo or otherwise relocate it before the Boltzmann tree release
   */
 object AccuracyTest extends SeedRandomMixIn {
 
-  val trainingDataFull: Seq[(Vector[Any], Double)] = TestUtils
-    .binTrainingData(
-      TestUtils.generateTrainingData(2048, 48, rng = rng),
-      inputBins = Seq((2, 32)) // bin the 3rd feature into a categorical
-    )
-    .asInstanceOf[Seq[(Vector[Any], Double)]]
+  val trainingDataFull: Seq[TrainingRow[Double]] =
+    DataGenerator.generate(2048, 48, rng = rng).withBinnedInputs(bins = Seq((2, 32))).data
 
   /**
     * Compute the RMSE and standard residual for a Boltzmann tree with the given temperature
@@ -107,7 +97,7 @@ object AccuracyTest extends SeedRandomMixIn {
       minInstances: Int,
       temperature: Double
   ): (Double, Double) = {
-    val trainingData = trainingDataFull.take(nRow).map { case (f, l) => (f.take(nFeat), l) }
+    val trainingData = trainingDataFull.take(nRow).map(_.mapInputs(inputs => inputs.take(nFeat)))
     val splitter = if (temperature > 0) {
       BoltzmannSplitter(temperature = temperature)
     } else {
@@ -118,16 +108,16 @@ object AccuracyTest extends SeedRandomMixIn {
       splitter = splitter,
       minLeafInstances = minInstances
     )
-    val learner = new Bagger(baseLearner, numBags = nRow * nScal, biasLearner = None)
-    val model = learner.train(trainingData, rng = rng).getModel()
-    val (features, labels) = trainingData.unzip
+    val learner = RegressionBagger(baseLearner, numBags = nRow * nScal, biasLearner = None)
+    val model = learner.train(trainingData, rng = rng).model
+    val (features, labels) = trainingData.map(row => (row.inputs, row.label)).unzip
     val predictions = model.transform(features)
-    val expected = predictions.getExpected().asInstanceOf[Seq[Double]]
-    val sigma = predictions.getUncertainty().get.asInstanceOf[Seq[Double]]
+    val expected = predictions.expected
+    val sigma = predictions.uncertainty().get.asInstanceOf[Seq[Double]]
     val pva: Seq[(Double, Double, Double)] = labels.indices.map { i =>
       (labels(i), expected(i), sigma(i))
     }
-    val rmse = Math.sqrt(pva.map { case (x, y, s) => Math.pow(x - y, 2) }.sum / pva.size)
+    val rmse = Math.sqrt(pva.map { case (x, y, _) => Math.pow(x - y, 2) }.sum / pva.size)
     val stdres = Math.sqrt(pva.map { case (x, y, s) => Math.pow((x - y) / s, 2) }.sum / pva.size)
     (rmse, stdres)
   }
diff --git a/src/test/scala/io/citrine/lolo/DataGenerator.scala b/src/test/scala/io/citrine/lolo/DataGenerator.scala
new file mode 100644
index 00000000..b2e183f2
--- /dev/null
+++ b/src/test/scala/io/citrine/lolo/DataGenerator.scala
@@ -0,0 +1,115 @@
+package io.citrine.lolo
+
+import breeze.stats.variance
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.linear.LinearRegressionLearner
+import io.citrine.lolo.stats.functions.Friedman
+import io.citrine.random.Random
+
+object DataGenerator {
+
+  /** Sequence of inputs and labels generated for training a learner. */
+  case class TrainingData[T](data: Seq[TrainingRow[T]]) {
+    def withBinnedInputs(bins: Seq[(Int, Int)]): TrainingData[T] = {
+      val baseInputs = data.map(_.inputs)
+      val indexToBins = bins.toMap
+      val transposedInputs = baseInputs.transpose
+      val binnedInputs = Vector
+        .tabulate(transposedInputs.length) { index =>
+          val indexData = transposedInputs(index).asInstanceOf[Seq[Double]]
+          indexToBins
+            .get(index)
+            .map { nBins => binData(indexData, nBins) }
+            .getOrElse(indexData)
+        }
+        .transpose
+      val binnedData = data.zip(binnedInputs).map { case (row, inputs) => row.withInputs(inputs) }
+      TrainingData(binnedData)
+    }
+  }
+
+  // An implicit class is used to restrict binning to when the labels are of type Double
+  implicit class TrainingLabelBinner(trainingData: TrainingData[Double]) {
+    def withBinnedLabels(bins: Int): TrainingData[String] = {
+      val baseLabels = trainingData.data.map(_.label)
+      val binnedLabels = binData(baseLabels, bins)
+      val binnedData = trainingData.data.zip(binnedLabels).map { case (row, label) => row.withLabel(label) }
+      TrainingData(binnedData)
+    }
+  }
+
+  def binData(continuousData: Seq[Double], nBins: Int): Seq[String] = {
+    val min = continuousData.min
+    val max = continuousData.max
+    continuousData.map(x => math.round(x * nBins / (max - min)).toString)
+  }
+
+  def generate(
+      rows: Int,
+      cols: Int,
+      function: Seq[Double] => Double = Friedman.friedmanGrosseSilverman,
+      xscale: Double = 1.0,
+      xoff: Double = 0.0,
+      noise: Double = 0.0,
+      rng: Random = Random()
+  ): TrainingData[Double] = {
+    val data = Vector.fill(rows) {
+      val input = Vector.fill(cols)(xscale * rng.nextDouble() + xoff)
+      TrainingRow(input, function(input) + noise * rng.nextGaussian())
+    }
+    TrainingData(data)
+  }
+
+  def iterate(
+      cols: Int,
+      function: Seq[Double] => Double = Friedman.friedmanGrosseSilverman,
+      xscale: Double = 1.0,
+      xoff: Double = 0.0,
+      noise: Double = 0.0,
+      rng: Random = Random()
+  ): Iterator[TrainingRow[Double]] = {
+    Iterator.continually {
+      val input = Vector.fill(cols)(xscale * rng.nextDouble() + xoff)
+      TrainingRow(input, function(input) + noise * rng.nextGaussian())
+    }
+  }
+
+  /**
+    * Enumerate the cartesian product of items in baseGrids.
+    *
+    * @param baseGrids a sequence of 1-d mesh specifications, one for each dimension of the output vectors
+    * @return a sequence of vectors enumerating the cartesian product of items in baseGrids
+    */
+  def enumerateGrid(baseGrids: Seq[Seq[Double]]): Seq[Vector[Double]] = {
+    if (baseGrids.length == 1) {
+      baseGrids.head.map { x => Vector(x) }
+    } else {
+      baseGrids.head.flatMap { x =>
+        enumerateGrid(baseGrids.takeRight(baseGrids.length - 1)).map { n => x +: n }
+      }
+    }
+  }
+
+  /**
+    * 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 makeLinearlyCorrelatedData(X: Seq[Double], rho: Double, rng: Random = 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) => TrainingRow(Vector(x), y) }).model
+    val yPred = linearModel.transform(X.map(Vector(_))).expected
+    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 }
+  }
+}
diff --git a/src/test/scala/io/citrine/lolo/PerformanceTest.scala b/src/test/scala/io/citrine/lolo/PerformanceTest.scala
index 0c12803a..ad8e9310 100644
--- a/src/test/scala/io/citrine/lolo/PerformanceTest.scala
+++ b/src/test/scala/io/citrine/lolo/PerformanceTest.scala
@@ -1,6 +1,8 @@
 package io.citrine.lolo
 
-import io.citrine.lolo.bags.{Bagger, MultiTaskBagger}
+import io.citrine.lolo.DataGenerator.TrainingData
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.bags.{Bagger, ClassificationBagger, MultiTaskBagger, RegressionBagger}
 import io.citrine.lolo.trees.classification.ClassificationTreeLearner
 import io.citrine.lolo.trees.multitask.MultiTaskTreeLearner
 import io.citrine.lolo.trees.regression.RegressionTreeLearner
@@ -17,46 +19,25 @@ import io.citrine.theta.Stopwatch
   */
 class PerformanceTest extends SeedRandomMixIn {
 
-  /**
-    * Time training and application of models
-    *
-    * @param trainingData which is used both to train and then later apply the models to
-    * @param n            number of training rows to take
-    * @param k            number of features to consider per split
-    * @param b            number of trees in the forest
-    * @param quiet        whether to print messages to the screen
-    * @return the training and application time, in seconds
-    */
-  def timedTest(
-      trainingData: Seq[(Vector[Any], Any)],
-      n: Int,
-      k: Int,
-      b: Int,
-      quiet: Boolean = true
+  private def timedTest[T](
+      bagger: Bagger[T],
+      trainingData: Seq[TrainingRow[T]]
   ): (Double, Double) = {
-    val data = trainingData.map(p => (p._1.take(k), p._2)).take(n)
-    val inputs = data.map(_._1)
-    val DTLearner = if (trainingData.head._2.isInstanceOf[Double]) {
-      new RegressionTreeLearner(numFeatures = k / 4)
-    } else {
-      new ClassificationTreeLearner(numFeatures = k / 4)
-    }
-    val baggedLearner = new Bagger(DTLearner, numBags = b)
-
+    val inputs = trainingData.map(_.inputs)
     val timeTraining = Stopwatch.time(
       {
-        baggedLearner.train(data).getModel()
+        bagger.train(trainingData).model
       },
       benchmark = "None",
       minRun = 4,
       targetError = 0.1,
       maxRun = 32
     )
-    val model = baggedLearner.train(data).getModel()
+    val model = bagger.train(trainingData).model
 
     val timePredicting = Stopwatch.time(
       {
-        model.transform(inputs).getUncertainty()
+        model.transform(inputs).uncertainty()
       },
       benchmark = "None",
       minRun = 4,
@@ -64,10 +45,33 @@ class PerformanceTest extends SeedRandomMixIn {
       maxRun = 32
     )
 
-    if (!quiet) println(f"${timeTraining}%10.4f, ${timePredicting}%10.4f, ${n}%6d, ${k}%6d, ${b}%6d")
     (timeTraining, timePredicting)
   }
 
+  def timeRegressor(
+      trainingData: Seq[TrainingRow[Double]],
+      n: Int,
+      k: Int,
+      b: Int
+  ): (Double, Double) = {
+    val data = trainingData.map(_.mapInputs(inputs => inputs.take(k))).take(n)
+    val baseLearner = RegressionTreeLearner(numFeatures = k / 4)
+    val bagger = RegressionBagger(baseLearner, numBags = b)
+    timedTest(bagger, data)
+  }
+
+  def timeClassifier[T](
+      trainingData: Seq[TrainingRow[T]],
+      n: Int,
+      k: Int,
+      b: Int
+  ): (Double, Double) = {
+    val data = trainingData.map(_.mapInputs(inputs => inputs.take(k))).take(n)
+    val baseLearner = ClassificationTreeLearner(numFeatures = k / 4)
+    val bagger = ClassificationBagger(baseLearner, numBags = b)
+    timedTest(bagger, data)
+  }
+
   // @Test
   def testScaling(): Unit = {
     val quiet: Boolean = true
@@ -75,12 +79,12 @@ class PerformanceTest extends SeedRandomMixIn {
     val Ks = Seq(8, 16, 32)
     val Bs = Seq(1024, 2048, 4096)
     if (!quiet) println(f"${"Train"}%10s, ${"Apply"}%10s, ${"N"}%6s, ${"K"}%6s, ${"B"}%6s")
-    timedTest(trainingData, Ns.head, Ks.head, Bs.head, true)
-    val (bTrain, bApply) = Bs.map(b => timedTest(trainingData, Ns.head, Ks.head, b, quiet)).unzip
+    timeRegressor(regressionData.data, Ns.head, Ks.head, Bs.head)
+    val (bTrain, bApply) = Bs.map(b => timeRegressor(regressionData.data, Ns.head, Ks.head, b)).unzip
     val (kTrain, kApply) =
-      (bTrain.zip(bApply).take(1) ++ Ks.tail.map(k => timedTest(trainingData, Ns.head, k, Bs.head, quiet))).unzip
+      (bTrain.zip(bApply).take(1) ++ Ks.tail.map(k => timeRegressor(regressionData.data, Ns.head, k, Bs.head))).unzip
     val (nTrain, nApply) =
-      (bTrain.zip(bApply).take(1) ++ Ns.tail.map(n => timedTest(trainingData, n, Ks.head, Bs.head, quiet))).unzip
+      (bTrain.zip(bApply).take(1) ++ Ns.tail.map(n => timeRegressor(regressionData.data, n, Ks.head, Bs.head))).unzip
 
     val bTrainScale = (1 until bTrain.size).map(i => bTrain(i) / bTrain(i - 1))
     val nTrainScale = (1 until nTrain.size).map(i => nTrain(i) / nTrain(i - 1))
@@ -104,20 +108,25 @@ class PerformanceTest extends SeedRandomMixIn {
     * Test the absolute performance to check for overall regressions
     */
   def testAbsolute(): Unit = {
-    val (nominalTrain, nominalPredict) = timedTest(classificationData, 1024, 32, 1024)
+    val (nominalTrain, nominalPredict) = timeClassifier(classificationData.data, 1024, 32, 1024)
     println(nominalTrain, nominalPredict)
   }
 
   def testMultitaskOverhead(N: Int): (Double, Double) = {
-    val subset = trainingData.take(N)
-    val (inputs, realLabels) = subset.unzip
-    val catLabels = realLabels.map(_ > realLabels.sum / realLabels.size)
-    val labels = Vector(realLabels, catLabels).transpose
+    val realRows = regressionData.data.take(N)
+    val avgReal = realRows.map(_.label).sum / realRows.length
+
+    val catLabels = realRows.map(_.label > avgReal)
+    val catRows = realRows.zip(catLabels).map { case (row, cat) => row.withLabel(cat) }
+
+    val multiRows = realRows.zip(catRows).map {
+      case (real, cat) => real.withLabel(Vector(real.label, cat.label))
+    }
 
     val trainSingle: Double = Stopwatch.time(
       {
-        new Bagger(RegressionTreeLearner()).train(inputs.zip(realLabels)).getLoss()
-        new Bagger(ClassificationTreeLearner()).train(inputs.zip(catLabels)).getLoss()
+        RegressionBagger(RegressionTreeLearner()).train(realRows).loss
+        ClassificationBagger(ClassificationTreeLearner()).train(catRows).loss
       },
       minRun = 1,
       maxRun = 1
@@ -125,7 +134,7 @@ class PerformanceTest extends SeedRandomMixIn {
 
     val trainMulti: Double = Stopwatch.time(
       {
-        MultiTaskBagger(MultiTaskTreeLearner()).train(inputs.zip(labels)).getLoss()
+        MultiTaskBagger(MultiTaskTreeLearner()).train(multiRows, rng = rng).loss
       },
       minRun = 1,
       maxRun = 1
@@ -134,13 +143,12 @@ class PerformanceTest extends SeedRandomMixIn {
     (trainMulti, trainSingle)
   }
 
-  val trainingData = TestUtils.generateTrainingData(2048, 37, rng = rng)
-  val classificationData = TestUtils.binTrainingData(trainingData, responseBins = Some(8))
+  val regressionData: TrainingData[Double] = DataGenerator.generate(2048, 37, rng = rng)
+  val classificationData: TrainingData[String] = regressionData.withBinnedLabels(bins = 8)
 }
 
 object PerformanceTest {
   def main(argv: Array[String]): Unit = {
     new PerformanceTest().testAbsolute()
-
   }
 }
diff --git a/src/test/scala/io/citrine/lolo/TestUtils.scala b/src/test/scala/io/citrine/lolo/TestUtils.scala
index c7abc0bd..d1189774 100644
--- a/src/test/scala/io/citrine/lolo/TestUtils.scala
+++ b/src/test/scala/io/citrine/lolo/TestUtils.scala
@@ -1,11 +1,6 @@
 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 io.citrine.random.Random
-import org.apache.commons.math3.random.MersenneTwister
 import org.junit.Before
 
 import scala.util.Try
@@ -40,98 +35,6 @@ object TestUtils {
     bs.close()
     res
   }
-
-  def generateTrainingData(
-      rows: Int,
-      cols: Int,
-      function: (Seq[Double] => Double) = Friedman.friedmanGrosseSilverman,
-      xscale: Double = 1.0,
-      xoff: Double = 0.0,
-      noise: Double = 0.0,
-      rng: Random = Random()
-  ): Vector[(Vector[Double], Double)] = {
-    Vector.fill(rows) {
-      val input = Vector.fill(cols)(xscale * rng.nextDouble() + xoff)
-      (input, function(input) + noise * rng.nextGaussian())
-    }
-  }
-
-  /** 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 = 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,
-      xscale: Double = 1.0,
-      xoff: Double = 0.0,
-      noise: Double = 0.0,
-      rng: Random = Random()
-  ): Iterator[(Vector[Double], Double)] = {
-    Iterator.continually {
-      val input = Vector.fill(cols)(xscale * rng.nextDouble() + xoff)
-      (input, function(input) + noise * rng.nextGaussian())
-    }
-  }
-
-  def binTrainingData(
-      continuousData: Seq[(Vector[Double], Double)],
-      inputBins: Seq[(Int, Int)] = Seq(),
-      responseBins: Option[Int] = None
-  ): Seq[(Vector[Any], Any)] = {
-    var outputData: Seq[(Vector[Any], Any)] = continuousData
-    inputBins.foreach {
-      case (index, nBins) =>
-        outputData = outputData.map {
-          case (input, response) =>
-            (input.updated(index, Math.round(input(index).asInstanceOf[Double] * nBins).toString), response)
-        }
-    }
-    responseBins.foreach { nBins =>
-      val max = continuousData.map(_._2).max
-      val min = continuousData.map(_._2).min
-      outputData = outputData.map {
-        case (input, response) =>
-          (input, Math.round(response.asInstanceOf[Double] * nBins / (max - min)).toString)
-      }
-    }
-    outputData
-  }
-
-  /**
-    * Enumerate the cartesian product of items in baseGrids.
-    *
-    * @param baseGrids a sequence of 1-d mesh specifications, one for each dimension of the output vectors
-    * @return a sequence of vectors enumerating the cartesian product of items in baseGrids
-    */
-  def enumerateGrid(baseGrids: Seq[Seq[Double]]): Seq[Vector[Double]] = {
-    if (baseGrids.length == 1) {
-      baseGrids.head.map { x => Vector(x) }
-    } else {
-      baseGrids.head.flatMap { x =>
-        enumerateGrid(baseGrids.takeRight(baseGrids.length - 1)).map { n => x +: n }
-      }
-    }
-  }
-
 }
 
 /**
diff --git a/src/test/scala/io/citrine/lolo/bags/BaggedResultTest.scala b/src/test/scala/io/citrine/lolo/bags/BaggedPredictionTest.scala
similarity index 76%
rename from src/test/scala/io/citrine/lolo/bags/BaggedResultTest.scala
rename to src/test/scala/io/citrine/lolo/bags/BaggedPredictionTest.scala
index 4723b277..0220cbfb 100644
--- a/src/test/scala/io/citrine/lolo/bags/BaggedResultTest.scala
+++ b/src/test/scala/io/citrine/lolo/bags/BaggedPredictionTest.scala
@@ -1,44 +1,57 @@
 package io.citrine.lolo.bags
 
 import breeze.stats.distributions.{Beta, RandBasis}
-import io.citrine.lolo.{RegressionResult, SeedRandomMixIn, TestUtils}
+import io.citrine.lolo.api.{RegressionResult, TrainingRow}
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
 import io.citrine.lolo.linear.GuessTheMeanLearner
 import io.citrine.lolo.stats.functions.Friedman
 import io.citrine.lolo.trees.regression.RegressionTreeLearner
 import org.junit.Test
 import org.scalatest.Assertions._
 
-class BaggedResultTest extends SeedRandomMixIn {
+class BaggedPredictionTest extends SeedRandomMixIn {
 
   @Test
   def testSingleMultiConsistency(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(512, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      inputBins = Seq((0, 8))
-    )
+    val trainingData = DataGenerator
+      .generate(512, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .data
 
     val DTLearner = RegressionTreeLearner(numFeatures = 12)
     val biasLearner = RegressionTreeLearner(maxDepth = 5, leafLearner = Some(GuessTheMeanLearner()))
 
     Array(
-      Bagger(DTLearner, numBags = 64, biasLearner = None, uncertaintyCalibration = false, useJackknife = true),
-      Bagger(
+      RegressionBagger(
+        DTLearner,
+        numBags = 64,
+        biasLearner = None,
+        uncertaintyCalibration = false,
+        useJackknife = true
+      ),
+      RegressionBagger(
         DTLearner,
         numBags = 64,
         biasLearner = Some(biasLearner),
         uncertaintyCalibration = true,
         useJackknife = false
       ),
-      Bagger(
+      RegressionBagger(
         DTLearner,
         numBags = 64,
         biasLearner = Some(biasLearner),
         uncertaintyCalibration = true,
         useJackknife = true
       ),
-      Bagger(DTLearner, numBags = 64, biasLearner = None, uncertaintyCalibration = false, useJackknife = false)
+      RegressionBagger(
+        DTLearner,
+        numBags = 64,
+        biasLearner = None,
+        uncertaintyCalibration = false,
+        useJackknife = false
+      )
     ).foreach { bagger =>
-      testConsistency(trainingData, bagger.train(trainingData, rng = rng).getModel())
+      testConsistency(trainingData, bagger.train(trainingData, rng = rng).model)
     }
   }
 
@@ -60,21 +73,23 @@ class BaggedResultTest extends SeedRandomMixIn {
 
             val sigmaObsAndSigmaMean: Seq[(Double, Double)] = (1 to 20).flatMap { _ =>
               val trainingDataTmp =
-                TestUtils.generateTrainingData(nRows, nCols, noise = 0.0, function = _ => 0.0, rng = rng)
-              val trainingData = trainingDataTmp.map { x => (x._1, x._2 + noiseLevel * rng.nextDouble()) }
-              val baggedLearner = Bagger(baseLearner, numBags = nBags, uncertaintyCalibration = true)
+                DataGenerator.generate(nRows, nCols, noise = 0.0, function = _ => 0.0, rng = rng).data
+              val trainingData = trainingDataTmp.map { row =>
+                row.mapLabel(_ + noiseLevel * rng.nextDouble())
+              }
+              val baggedLearner = RegressionBagger(baseLearner, numBags = nBags, uncertaintyCalibration = true)
               val RFMeta = baggedLearner.train(trainingData, rng = rng)
-              val RF = RFMeta.getModel()
-              val results = RF.transform(trainingData.take(4).map(_._1))
+              val RF = RFMeta.model
+              val results = RF.transform(trainingData.take(4).map(_.inputs))
 
-              val sigmaMean: Seq[Double] = results.getUncertainty(observational = false).get.asInstanceOf[Seq[Double]]
-              sigmaMean.zip(results.asInstanceOf[RegressionResult].getStdDevMean().get).foreach {
+              val sigmaMean: Seq[Double] = results.uncertainty(observational = false).get.asInstanceOf[Seq[Double]]
+              sigmaMean.zip(results.asInstanceOf[RegressionResult].stdDevMean.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[RegressionResult].getStdDevObs().get).foreach {
+              val sigmaObs: Seq[Double] = results.uncertainty().get.asInstanceOf[Seq[Double]]
+              sigmaObs.zip(results.asInstanceOf[RegressionResult].stdDevObs.get).foreach {
                 case (a, b) =>
                   assert(a == b, s"Expected getUncertainty()=getStdDevObs() for $configDescription")
               }
@@ -155,24 +170,24 @@ class BaggedResultTest extends SeedRandomMixIn {
     * @param trainingData The original training data for the model
     * @param model        The trained model
     */
-  private def testConsistency(trainingData: Seq[(Vector[Any], Any)], model: BaggedModel[Any]): Unit = {
+  private def testConsistency(trainingData: Seq[TrainingRow[Double]], model: BaggedModel[Double]): Unit = {
     val testSubset = rng.shuffle(trainingData).take(16)
     val (singleValues, singleObsUnc, singleMeanUnc) = testSubset.map {
-      case (x, _) =>
+      case TrainingRow(x, _, _) =>
         val res = model.transform(Seq(x))
         (
-          res.getExpected().head.asInstanceOf[Double],
-          res.getUncertainty(true).get.head.asInstanceOf[Double],
-          res.getUncertainty(false).get.head.asInstanceOf[Double]
+          res.expected.head,
+          res.uncertainty(true).get.head.asInstanceOf[Double],
+          res.uncertainty(false).get.head.asInstanceOf[Double]
         )
     }.unzip3
 
     val (multiValues, multiObsUnc, multiMeanUnc) = {
-      val res = model.transform(testSubset.map(_._1))
+      val res = model.transform(testSubset.map(_.inputs))
       (
-        res.getExpected().map(_.asInstanceOf[Double]),
-        res.getUncertainty(true).get.map(_.asInstanceOf[Double]),
-        res.getUncertainty(false).get.map(_.asInstanceOf[Double])
+        res.expected,
+        res.uncertainty(true).get.map(_.asInstanceOf[Double]),
+        res.uncertainty(false).get.map(_.asInstanceOf[Double])
       )
     }
 
diff --git a/src/test/scala/io/citrine/lolo/bags/BaggerTest.scala b/src/test/scala/io/citrine/lolo/bags/BaggerTest.scala
index deecb2c9..86d21012 100644
--- a/src/test/scala/io/citrine/lolo/bags/BaggerTest.scala
+++ b/src/test/scala/io/citrine/lolo/bags/BaggerTest.scala
@@ -1,10 +1,12 @@
 package io.citrine.lolo.bags
 
 import breeze.linalg.DenseMatrix
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn, TestUtils}
 import io.citrine.lolo.linear.{GuessTheMeanLearner, LinearRegressionLearner}
 import io.citrine.lolo.stats.functions.Friedman
-import io.citrine.lolo.transformers.{FeatureRotator, Standardizer}
+import io.citrine.lolo.transformers.rotator.FeatureRotator
+import io.citrine.lolo.transformers.standardizer.RegressionStandardizer
 import io.citrine.lolo.trees.classification.ClassificationTreeLearner
 import io.citrine.lolo.trees.regression.RegressionTreeLearner
 import io.citrine.lolo.trees.splits.RegressionSplitter
@@ -12,11 +14,7 @@ import org.junit.Test
 import org.scalatest.Assertions._
 
 import java.io.{File, PrintWriter}
-import java.util.concurrent._
 
-/**
-  * Created by maxhutch on 11/29/16.
-  */
 @Test
 class BaggerTest extends SeedRandomMixIn {
 
@@ -30,26 +28,22 @@ class BaggerTest extends SeedRandomMixIn {
       x.zip(beta).map { case (xi, w) => xi * w }.sum
     }
 
-    val trainingData = TestUtils.generateTrainingData(
-      rows = 256,
-      beta0.length,
-      noise = 0.5,
-      function = x => linearFunction(x, beta0),
-      rng = rng
-    )
+    val trainingData = DataGenerator
+      .generate(rows = 256, beta0.length, noise = 0.5, function = x => linearFunction(x, beta0), rng = rng)
+      .data
 
-    val baseLearner = new Standardizer(LinearRegressionLearner(regParam = Some(0.5)))
-    val baggedLearner = Bagger(baseLearner, numBags = trainingData.size)
+    val baseLearner = RegressionStandardizer(LinearRegressionLearner(regParam = Some(0.5)))
+    val baggedLearner = RegressionBagger(baseLearner, numBags = trainingData.size)
 
     val learnerMeta = baggedLearner.train(trainingData, rng = rng)
-    val model = learnerMeta.getModel()
+    val model = learnerMeta.model
 
-    assert(learnerMeta.getLoss().get < 1.0, "Loss of bagger is larger than expected")
+    assert(learnerMeta.loss.get < 1.0, "Loss of bagger is larger than expected")
 
-    val prediction = model.transform(trainingData.take(1).map(_._1))
-    val sigma = prediction.getUncertainty().get.map(_.asInstanceOf[Double])
+    val prediction = model.transform(trainingData.take(1).map(_.inputs))
+    val sigma = prediction.uncertainty().get.map(_.asInstanceOf[Double])
     assert(sigma.forall(_ > 0.0))
-    assert(prediction.getGradient().isDefined, "No gradient returned for linear ensemble.")
+    assert(prediction.gradient.isDefined, "No gradient returned for linear ensemble.")
   }
 
   /**
@@ -57,27 +51,26 @@ class BaggerTest extends SeedRandomMixIn {
     */
   @Test
   def testRegressionBagger(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      inputBins = Seq((0, 8))
-    )
+    val trainingData = DataGenerator
+      .generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .data
+
     val DTLearner = RegressionTreeLearner(numFeatures = 3)
-    val baggedLearner =
-      Bagger(DTLearner, numBags = trainingData.size)
+    val baggedLearner = RegressionBagger(DTLearner, numBags = trainingData.size)
     val RFMeta = baggedLearner.train(trainingData, rng = rng)
-    val RF = RFMeta.getModel()
+    val RF = RFMeta.model
 
-    assert(RFMeta.getLoss().get < 1.0, "Loss of bagger is larger than expected")
+    assert(RFMeta.loss.get < 1.0, "Loss of bagger is larger than expected")
 
-    val results = RF.transform(trainingData.map(_._1))
-    val means = results.getExpected()
-    val sigma: Seq[Double] = results.getUncertainty().get.asInstanceOf[Seq[Double]]
+    val results = RF.transform(trainingData.map(_.inputs))
+    val sigma: Seq[Double] = results.uncertainty().get.asInstanceOf[Seq[Double]]
     assert(sigma.forall(_ > 0.0))
 
-    assert(results.getGradient().isEmpty, "Returned a gradient when there shouldn't be one")
+    assert(results.gradient.isEmpty, "Returned a gradient when there shouldn't be one")
 
     /* The first feature should be the most important */
-    val importances = RFMeta.getFeatureImportance().get
+    val importances = RFMeta.featureImportance.get
     assert(importances(1) == importances.max)
   }
 
@@ -86,34 +79,33 @@ class BaggerTest extends SeedRandomMixIn {
     */
   @Test
   def testClassificationBagger(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      inputBins = Seq((0, 8)),
-      responseBins = Some(8)
-    )
+    val trainingData = DataGenerator
+      .generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .withBinnedLabels(bins = 8)
+      .data
+
     val DTLearner = ClassificationTreeLearner()
-    val baggedLearner =
-      Bagger(DTLearner, numBags = trainingData.size / 2)
+    val baggedLearner = ClassificationBagger(DTLearner, numBags = trainingData.size / 2)
     val RFMeta = baggedLearner.train(trainingData, rng = rng)
-    val RF = RFMeta.getModel()
+    val RF = RFMeta.model
 
     /* Inspect the results */
-    val results = RF.transform(trainingData.map(_._1))
-    val means = results.getExpected()
-    assert(trainingData.map(_._2).zip(means).forall { case (a, p) => a == p })
+    val results = RF.transform(trainingData.map(_.inputs))
+    val means = results.expected
+    assert(trainingData.map(_.label).zip(means).forall { case (a, p) => a == p })
 
-    val uncertainty = results.getUncertainty()
+    val uncertainty = results.uncertainty()
     assert(uncertainty.isDefined)
-    assert(trainingData.map(_._2).zip(uncertainty.get).forall {
+    assert(trainingData.map(_.label).zip(uncertainty.get).forall {
       case (a, probs) =>
-        val classProbabilities = probs.asInstanceOf[Map[Any, Double]]
-        val maxProb = classProbabilities(a)
-        maxProb > 0.5 && maxProb < 1.0 && Math.abs(classProbabilities.values.sum - 1.0) < 1.0e-6
+        val maxProb = probs(a)
+        maxProb > 0.5 && maxProb < 1.0 && Math.abs(probs.values.sum - 1.0) < 1.0e-6
     })
-    assert(results.getGradient().isEmpty, "Returned a gradient when there shouldn't be one")
+    assert(results.gradient.isEmpty, "Returned a gradient when there shouldn't be one")
 
     /* The first feature should be the most important */
-    val importances = RFMeta.getFeatureImportance().get
+    val importances = RFMeta.featureImportance.get
     assert(importances.slice(0, 5).min > importances.slice(5, importances.size).max)
   }
 
@@ -129,20 +121,19 @@ class BaggerTest extends SeedRandomMixIn {
     val width = 0.10 // make the function more linear
     val nFeatures = 5
     val bagsPerRow = 4 // picked to be large enough that bias correction is small but model isn't too expensive
-    val trainingData = TestUtils.generateTrainingData(128, nFeatures, xscale = width, rng = rng)
+    val trainingData = DataGenerator.generate(128, nFeatures, xscale = width, rng = rng).data
     val DTLearner = RegressionTreeLearner(numFeatures = nFeatures)
     val bias = RegressionTreeLearner(maxDepth = 4)
-    val baggedLearner = Bagger(
-      DTLearner,
-      numBags = bagsPerRow * trainingData.size,
-      biasLearner = Some(bias)
-    )
+    val baggedLearner = RegressionBagger(DTLearner, numBags = bagsPerRow * trainingData.size, biasLearner = Some(bias))
     val RFMeta = baggedLearner.train(trainingData, rng = rng)
-    val RF = RFMeta.getModel()
+    val RF = RFMeta.model
 
-    val interiorTestSet =
-      TestUtils.generateTrainingData(128, nFeatures, xscale = width / 2.0, xoff = width / 4.0, rng = rng)
-    val fullTestSet = TestUtils.generateTrainingData(128, nFeatures, xscale = width, rng = rng)
+    val interiorTestSet = DataGenerator
+      .generate(128, nFeatures, xscale = width / 2.0, xoff = width / 4.0, rng = rng)
+      .data
+      .map(row => (row.inputs, row.label))
+    val fullTestSet =
+      DataGenerator.generate(128, nFeatures, xscale = width, rng = rng).data.map(row => (row.inputs, row.label))
 
     val interiorStandardRMSE = BaggerTest.getStandardRMSE(interiorTestSet, RF)
     val fullStandardRMSE = BaggerTest.getStandardRMSE(fullTestSet, RF)
@@ -160,9 +151,7 @@ class BaggerTest extends SeedRandomMixIn {
   def testUncertaintyCalibrationWithConstantResponse(): Unit = {
     // setup some training data with constant labels
     val nFeatures = 5
-    val X: Vector[Vector[Any]] =
-      TestUtils.generateTrainingData(128, nFeatures, xscale = 0.5, rng = rng).map(_._1)
-    val y: Vector[Any] = X.map(_ => 0.0)
+    val constantData = DataGenerator.generate(128, nFeatures, xscale = 0.5, rng = rng).data.map(_.withLabel(0.0))
 
     // setup a relatively complicated random forest (turn a bunch of stuff on)
     val DTLearner = RegressionTreeLearner(
@@ -172,8 +161,8 @@ class BaggerTest extends SeedRandomMixIn {
       splitter = RegressionSplitter(randomizePivotLocation = true)
     )
 
-    val bagger = new Bagger(
-      new Standardizer(DTLearner),
+    val bagger = RegressionBagger(
+      RegressionStandardizer(DTLearner),
       numBags = 64,
       useJackknife = true,
       biasLearner = Some(
@@ -187,14 +176,13 @@ class BaggerTest extends SeedRandomMixIn {
     )
 
     // Make sure the model trains
-    val model = bagger.train(X.zip(y), rng = rng).getModel()
+    val model = bagger.train(constantData, rng = rng).model
 
     // Generate a new test set and make sure the predictions are 0 +/- 0
-    val testX: Vector[Vector[Any]] =
-      TestUtils.generateTrainingData(128, nFeatures, xscale = 0.5, rng = rng).map(_._1)
+    val testX = DataGenerator.generate(128, nFeatures, xscale = 0.5, rng = rng).data.map(_.inputs)
     val predictions = model.transform(testX)
-    assert(predictions.getExpected().forall(_ == 0.0))
-    assert(predictions.getUncertainty().get.forall(_ == 0.0))
+    assert(predictions.expected.forall(_ == 0.0))
+    assert(predictions.uncertainty().get.forall(_ == 0.0))
   }
 
   /**
@@ -208,17 +196,17 @@ class BaggerTest extends SeedRandomMixIn {
   @Test
   def testScores(): Unit = {
     val csv = TestUtils.readCsv("double_example.csv")
-    val trainingData = csv.map(vec => (vec.init, vec.last.asInstanceOf[Double]))
+    val trainingData = csv.map(vec => TrainingRow(vec.init, vec.last.asInstanceOf[Double]))
     val DTLearner = RegressionTreeLearner()
-    val baggedLearner = Bagger(
+    val baggedLearner = RegressionBagger(
       DTLearner,
       numBags = trainingData.size * 16
     ) // use lots of trees to reduce noise
-    val RF = baggedLearner.train(trainingData, rng = rng).getModel()
+    val RF = baggedLearner.train(trainingData, rng = rng).model
 
     /* Call transform on the training data */
-    val results = RF.transform(trainingData.map(_._1))
-    val scores = results.getImportanceScores().get
+    val results = RF.transform(trainingData.map(_.inputs))
+    val scores = results.importanceScores.get
     val corners = Seq(0, 7, 56, 63)
     corners.foreach { i =>
       assert(
@@ -233,29 +221,30 @@ class BaggerTest extends SeedRandomMixIn {
     */
   @Test
   def calibrationTimeTest(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      inputBins = Seq((0, 8))
-    )
+    val trainingData = DataGenerator
+      .generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .data
+
     val DTLearner = RegressionTreeLearner(numFeatures = 3)
     val start = System.nanoTime()
-    Bagger(
+    RegressionBagger(
       DTLearner,
       numBags = trainingData.size,
       uncertaintyCalibration = false
     )
       .train(trainingData, rng = rng)
-      .getModel()
+      .model
     val unCalibratedTime = 1.0e-9 * (System.nanoTime() - start)
 
     val startAgain = System.nanoTime()
-    Bagger(
+    RegressionBagger(
       DTLearner,
       numBags = 64,
       uncertaintyCalibration = true
     )
       .train(trainingData, rng = rng)
-      .getModel()
+      .model
     val calibratedTime = 1.0e-9 * (System.nanoTime() - startAgain)
 
     assert(calibratedTime < unCalibratedTime, s"The calibration scheme has experienced a dramatic slowdown")
@@ -273,7 +262,7 @@ class BaggerTest extends SeedRandomMixIn {
     // Define a simple, binary function and create training data
     def stepFunction(x: Seq[Double]): Double = Math.floor(2 * x(0))
 
-    val trainingData = TestUtils.generateTrainingData(rows = 16, cols = 2, function = stepFunction, rng = rng)
+    val trainingData = DataGenerator.generate(rows = 16, cols = 2, function = stepFunction, rng = rng).data
 
     /* Create a bagger out of GuessTheMean learners, and train the model.
      * This model has a rescale field, which should be a real number. If it is not,
@@ -284,17 +273,17 @@ class BaggerTest extends SeedRandomMixIn {
       numFeatures = 2,
       splitter = RegressionSplitter(randomizePivotLocation = true)
     )
-    val trainedModel: BaggedModel[Any] = Bagger(
+    val trainedModel = RegressionBagger(
       DTLearner,
       numBags = 16,
       useJackknife = true,
       uncertaintyCalibration = true
     )
       .train(trainingData, rng = rng)
-      .getModel()
+      .model
 
     try {
-      val _: BaggedModel[Any] = Bagger(
+      val _: BaggedModel[Any] = RegressionBagger(
         DTLearner,
         numBags = 16,
         useJackknife = true,
@@ -302,7 +291,7 @@ class BaggerTest extends SeedRandomMixIn {
         disableBootstrap = true
       )
         .train(trainingData, rng = rng)
-        .getModel()
+        .model
       fail("Setting both uncertaintyCalibration and disableBootstrap should throw an exception.")
     } catch {
       case _: Throwable =>
@@ -320,13 +309,13 @@ class BaggerTest extends SeedRandomMixIn {
   def testUncertaintyFloor(): Unit = {
     (0 until 16384).foreach { idx =>
       val trainingData =
-        TestUtils.generateTrainingData(16, 5, noise = 0.0, function = Friedman.friedmanSilverman, rng = rng)
+        DataGenerator.generate(16, 5, noise = 0.0, function = Friedman.friedmanSilverman, rng = rng).data
       val DTLearner = RegressionTreeLearner(numFeatures = 2)
-      val sigma = Bagger(DTLearner, numBags = 7)
+      val sigma = RegressionBagger(DTLearner, numBags = 7)
         .train(trainingData, rng = rng)
-        .getModel()
-        .transform(trainingData.map(_._1))
-        .getUncertainty()
+        .model
+        .transform(trainingData.map(_.inputs))
+        .uncertainty()
         .get
         .asInstanceOf[Seq[Double]]
       assert(sigma.forall(_ > 0.0), s"Found an predicted uncertainty of ${sigma.min} during trial $idx")
@@ -342,17 +331,17 @@ class BaggerTest extends SeedRandomMixIn {
   def testUncertaintyFloorWithBias(): Unit = {
     (0 until 1024).foreach { idx =>
       val trainingData =
-        TestUtils.generateTrainingData(16, 5, noise = 0.0, function = Friedman.friedmanSilverman, rng = rng)
+        DataGenerator.generate(16, 5, noise = 0.0, function = Friedman.friedmanSilverman, rng = rng).data
       val DTLearner = RegressionTreeLearner(numFeatures = 2)
-      val sigma = Bagger(
+      val sigma = RegressionBagger(
         DTLearner,
         numBags = 7,
         biasLearner = Some(GuessTheMeanLearner())
       )
         .train(trainingData, rng = rng)
-        .getModel()
-        .transform(trainingData.map(_._1))
-        .getUncertainty()
+        .model
+        .transform(trainingData.map(_.inputs))
+        .uncertainty()
         .get
         .asInstanceOf[Seq[Double]]
       assert(sigma.forall(_ > 0.0), s"Found an predicted uncertainty of ${sigma.min} during trial $idx")
@@ -365,20 +354,13 @@ class BaggerTest extends SeedRandomMixIn {
   @Test
   def testShapley(): Unit = {
     val nCols = 5
-    val trainingData = TestUtils.generateTrainingData(
-      64,
-      nCols,
-      noise = 0.0,
-      function = Friedman.friedmanSilverman,
-      rng = rng
-    )
+    val trainingData =
+      DataGenerator.generate(64, nCols, noise = 0.0, function = Friedman.friedmanSilverman, rng = rng).data
     val DTLearner = RegressionTreeLearner(numFeatures = nCols)
-    val model = Bagger(DTLearner)
-      .train(trainingData, rng = rng)
-      .getModel()
-    val trees = model.getModels()
+    val model = RegressionBagger(DTLearner).train(trainingData, rng = rng).model
+    val trees = model.ensembleModels
     trainingData.foreach {
-      case (x, _) =>
+      case TrainingRow(x, _, _) =>
         val shapley = model.shapley(x).get
 
         // Do a quick sanity check on the output format.
@@ -407,19 +389,14 @@ class BaggerTest extends SeedRandomMixIn {
   @Test
   def testShapleyIsEmpty(): Unit = {
     val nCols = 2
-    val trainingData = TestUtils.generateTrainingData(
-      8,
-      nCols,
-      noise = 0.0,
-      function = Friedman.friedmanSilverman,
-      rng = rng
-    )
+    val trainingData =
+      DataGenerator.generate(8, nCols, noise = 0.0, function = Friedman.friedmanSilverman, rng = rng).data
     val learner = FeatureRotator(RegressionTreeLearner(numFeatures = nCols))
-    val model = Bagger(learner)
+    val model = RegressionBagger(learner)
       .train(trainingData, rng = rng)
-      .getModel()
+      .model
 
-    val x = trainingData.head._1
+    val x = trainingData.head.inputs
     assert(model.shapley(x).isEmpty)
   }
 }
@@ -448,24 +425,19 @@ object BaggerTest extends SeedRandomMixIn {
       (4 to 8 by 2).foreach { nRowsLog: Int =>
         val nRows = 1 << nRowsLog
         (1 to 3).foreach { repNum =>
-          val trainingData = TestUtils.generateTrainingData(
-            nRows,
-            nCols,
-            noise = 0.0,
-            function = Friedman.friedmanSilverman,
-            rng = rng
-          )
+          val trainingData =
+            DataGenerator.generate(nRows, nCols, noise = 0.0, function = Friedman.friedmanSilverman, rng = rng).data
           val DTLearner = RegressionTreeLearner(numFeatures = nCols)
           println(s"Training model nCols=$nCols\tnRows=$nRows\trepNum=$repNum")
-          val model = Bagger(DTLearner)
+          val model = RegressionBagger(DTLearner)
             .train(trainingData, rng = rng)
-            .getModel()
+            .model
           println(s"Trained")
 
           rng.shuffle(trainingData).take(16).zipWithIndex.foreach {
             case (x, i) =>
               val t0 = System.nanoTime()
-              val shapley = model.shapley(x._1).get
+              val shapley = model.shapley(x.inputs).get
               val t1 = System.nanoTime()
               pw.write(s"$nCols\t$nRows\t$repNum\t$i\t${t1 - t0}\n")
               pw.flush()
@@ -481,11 +453,10 @@ object BaggerTest extends SeedRandomMixIn {
     val pva = testSet
       .map(_._2)
       .zip(
-        predictions
-          .getExpected()
+        predictions.expected
           .asInstanceOf[Seq[Double]]
           .zip(
-            predictions.getUncertainty().get.asInstanceOf[Seq[Double]]
+            predictions.uncertainty().get.asInstanceOf[Seq[Double]]
           )
       )
     val standardError = pva.map {
@@ -494,5 +465,4 @@ object BaggerTest extends SeedRandomMixIn {
     }
     Math.sqrt(standardError.map(Math.pow(_, 2.0)).sum / testSet.size)
   }
-
 }
diff --git a/src/test/scala/io/citrine/lolo/bags/MultiTaskBaggerTest.scala b/src/test/scala/io/citrine/lolo/bags/MultiTaskBaggerTest.scala
index 4115f097..73258b85 100644
--- a/src/test/scala/io/citrine/lolo/bags/MultiTaskBaggerTest.scala
+++ b/src/test/scala/io/citrine/lolo/bags/MultiTaskBaggerTest.scala
@@ -1,7 +1,8 @@
 package io.citrine.lolo.bags
 
 import breeze.stats.distributions.{Beta, RandBasis}
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
 import io.citrine.lolo.linear.GuessTheMeanLearner
 import io.citrine.lolo.stats.functions.Friedman
 import io.citrine.lolo.stats.metrics.ClassificationMetrics
@@ -11,9 +12,6 @@ import io.citrine.lolo.trees.regression.RegressionTreeLearner
 import org.junit.Test
 import org.scalatest.Assertions._
 
-/**
-  * Created by maxhutch on 11/29/16.
-  */
 @Test
 class MultiTaskBaggerTest extends SeedRandomMixIn {
 
@@ -22,11 +20,12 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
     */
   @Test
   def testSingleRegression(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(512, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      inputBins = Seq((0, 8))
-    )
-    val reshapedTrainingData = trainingData.map { case (input, label) => (input, Vector(label)) }
+    val trainingData = DataGenerator
+      .generate(512, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .data
+
+    val reshapedTrainingData = trainingData.map(_.mapLabel(Vector(_)))
     val DTLearner = MultiTaskTreeLearner()
     val baggedLearner = MultiTaskBagger(
       DTLearner,
@@ -34,18 +33,18 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
       uncertaintyCalibration = true
     )
     val RFMeta = baggedLearner.train(reshapedTrainingData, rng = rng)
-    val RF = RFMeta.getModels().head
+    val RF = RFMeta.models.head
 
-    val results = RF.transform(trainingData.map(_._1))
-    val means = results.getExpected()
-    val sigma: Seq[Double] = results.getUncertainty().get.asInstanceOf[Seq[Double]]
+    val results = RF.transform(trainingData.map(_.inputs))
+    val means = results.expected
+    val sigma: Seq[Double] = results.uncertainty().get.asInstanceOf[Seq[Double]]
     assert(sigma.forall(_ >= 0.0))
     assert(
-      results.asInstanceOf[MultiPredictionBaggedResult].rescale != 1.0,
+      results.asInstanceOf[MultiPointBaggedPrediction].rescaleRatio != 1.0,
       "uncertainty calibration ratio was not included in prediction result"
     )
-    assert(results.getGradient().isEmpty, "Returned a gradient when there shouldn't be one")
-    assert(RFMeta.getLoss().get < 1.0, "Loss of bagger is larger than expected")
+    assert(results.gradient.isEmpty, "Returned a gradient when there shouldn't be one")
+    assert(RFMeta.loss.get < 1.0, "Loss of bagger is larger than expected")
   }
 
   /**
@@ -65,27 +64,27 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
               s"learner=${baseLearner.getClass.toString}\tnRows=$nRows\tnCols=$nCols\tnumBags=$nBags"
 
             val sigmaObsAndSigmaMean: Seq[(Double, Double)] = (1 to 20).flatMap { _ =>
-              val trainingDataTmp =
-                TestUtils.generateTrainingData(nRows, nCols, function = _ => 0.0, rng = rng)
-              val trainingData = trainingDataTmp.map { x => (x._1, x._2 + noiseLevel * rng.nextDouble()) }
-              val reshapedTrainingData = trainingData.map { case (input, label) => (input, Vector(label)) }
+              val trainingDataTmp = DataGenerator.generate(nRows, nCols, function = _ => 0.0, rng = rng).data
+              val trainingData = trainingDataTmp.map(_.mapLabel(_ + noiseLevel * rng.nextDouble()))
+
+              val reshapedTrainingData = trainingData.map(_.mapLabel(Vector(_)))
               val baggedLearner = MultiTaskBagger(
                 baseLearner,
                 numBags = nBags,
                 uncertaintyCalibration = true
               )
               val RFMeta = baggedLearner.train(reshapedTrainingData, rng = rng)
-              val RF = RFMeta.getModels().head
-              val results = RF.transform(trainingData.take(4).map(_._1))
+              val RF = RFMeta.models.head
+              val results = RF.transform(trainingData.take(4).map(_.inputs))
 
-              val sigmaMean: Seq[Double] = results.getUncertainty(observational = false).get.asInstanceOf[Seq[Double]]
-              sigmaMean.zip(results.asInstanceOf[MultiPredictionBaggedResult].getStdDevMean().get).foreach {
+              val sigmaMean: Seq[Double] = results.uncertainty(observational = false).get.asInstanceOf[Seq[Double]]
+              sigmaMean.zip(results.asInstanceOf[MultiPointBaggedPrediction].stdDevMean.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[MultiPredictionBaggedResult].getStdDevObs().get).foreach {
+              val sigmaObs: Seq[Double] = results.uncertainty().get.asInstanceOf[Seq[Double]]
+              sigmaObs.zip(results.asInstanceOf[MultiPointBaggedPrediction].stdDevObs.get).foreach {
                 case (a, b) =>
                   assert(a == b, s"Expected getUncertainty()=getStdDevObs() for $configDescription")
               }
@@ -163,26 +162,27 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
     */
   @Test
   def testClassificationBagger(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(128, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      inputBins = Seq((0, 8)),
-      responseBins = Some(8)
-    )
-    val reshapedTrainingData = trainingData.map { case (input, label) => (input, Vector(label)) }
+    val trainingData = DataGenerator
+      .generate(128, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .withBinnedLabels(bins = 8)
+      .data
+
+    val reshapedTrainingData = trainingData.map(_.mapLabel(Vector(_)))
     val DTLearner = MultiTaskTreeLearner()
     val baggedLearner =
       MultiTaskBagger(DTLearner, numBags = trainingData.size)
     val RFMeta = baggedLearner.train(reshapedTrainingData, rng = rng)
-    val RF = RFMeta.getModels().head
+    val RF = RFMeta.models.head
 
     /* Inspect the results */
-    val results = RF.transform(trainingData.map(_._1))
-    val means = results.getExpected()
-    assert(trainingData.map(_._2).zip(means).forall { case (a, p) => a == p })
+    val results = RF.transform(trainingData.map(_.inputs))
+    val means = results.expected
+    assert(trainingData.map(_.label).zip(means).forall { case (a, p) => a == p })
 
-    val uncertainty = results.getUncertainty()
+    val uncertainty = results.uncertainty()
     assert(uncertainty.isDefined)
-    trainingData.map(_._2).zip(uncertainty.get).foreach {
+    trainingData.map(_.label).zip(uncertainty.get).foreach {
       case (a, probs) =>
         val classProbabilities = probs.asInstanceOf[Map[Any, Double]]
         val maxProb = classProbabilities(a)
@@ -190,7 +190,7 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
         assert(maxProb < 1.0)
         assert(Math.abs(classProbabilities.values.sum - 1.0) < 1.0e-6)
     }
-    assert(results.getGradient().isEmpty, "Returned a gradient when there shouldn't be one")
+    assert(results.gradient.isEmpty, "Returned a gradient when there shouldn't be one")
   }
 
   /** Test that a multi-task bagged model properly stores and transposes individual trees, and remembers labels. */
@@ -199,11 +199,14 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
     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, rng = rng)
-    val (inputs: Seq[Vector[Double]], realLabel: Seq[Double]) = raw.unzip
-    val catLabel: Seq[Boolean] = raw.map(_._2 > realLabel.max / 2.0)
+
+    val realRows = DataGenerator
+      .generate(numTrain, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .data
+    val (inputs, realLabel) = realRows.map(row => (row.inputs, row.label)).unzip
+    val catLabel = realLabel.map(_ > realLabel.max / 2.0)
     val labels = Vector(realLabel, catLabel).transpose
+    val multiTaskRows = TrainingRow.build(inputs.zip(labels))
 
     val learner = MultiTaskTreeLearner()
     val baggedLearner = MultiTaskBagger(
@@ -212,19 +215,19 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
       biasLearner = Some(RegressionTreeLearner(maxDepth = 2)),
       uncertaintyCalibration = true
     )
-    val RF = baggedLearner.train(inputs.zip(labels), rng = rng)
+    val RF = baggedLearner.train(multiTaskRows, rng = rng)
 
     val testInputs = inputs.take(numTest)
-    val predictionResult = RF.getModel().transform(testInputs)
-    assert(predictionResult.predictions.length == numBags)
+    val predictionResult = RF.model.transform(testInputs)
+    assert(predictionResult.ensemblePredictions.length == numBags)
 
     // because the uncertainty is recalibrated, the prediction result should have a rescale value that is not equal to 1.0
-    assert(predictionResult.baggedPredictions.head.asInstanceOf[MultiPredictionBaggedResult].rescale != 1.0)
+    assert(predictionResult.labelPredictions.head.asInstanceOf[MultiPointBaggedPrediction].rescaleRatio != 1.0)
 
     // 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()
+    val expected = predictionResult.expected
+    val expectedCat = RF.models(1).transform(testInputs).expected
     (0 until numTest).foreach { i =>
       assert(expected(i)(1) == catLabel(i))
       assert(catLabel(i) == expectedCat(i))
@@ -238,12 +241,15 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
     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, rng = rng)
-    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, rng = rng)
+
+    val realRows = DataGenerator
+      .generate(numTrain, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .data
+    val (inputs, realLabel) = realRows.map(row => (row.inputs, row.label)).unzip
+    val catLabel = realLabel.map(_ > realLabel.max / 2.0)
+    val correlatedLabel = DataGenerator.makeLinearlyCorrelatedData(realLabel, trainingRho, rng = rng)
     val labels = Vector(realLabel, catLabel, correlatedLabel).transpose
+    val multiTaskRows = TrainingRow.build(inputs.zip(labels))
 
     val learner = MultiTaskTreeLearner()
     val baggedLearner = MultiTaskBagger(
@@ -251,19 +257,19 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
       numBags = numBags,
       biasLearner = Some(RegressionTreeLearner(maxDepth = 2))
     )
-    val RF = baggedLearner.train(inputs.zip(labels), rng = rng).getModel()
+    val RF = baggedLearner.train(multiTaskRows, rng = rng).model
 
     val testInputs =
-      TestUtils.generateTrainingData(numTest, 12, function = Friedman.friedmanSilverman, rng = rng).map(_._1)
+      DataGenerator.generate(numTest, 12, function = Friedman.friedmanSilverman, rng = rng).data.map(_.inputs)
     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))
+      assert(predictionResult.uncertaintyCorrelation(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)
+      assert(predictionResult.uncertaintyCorrelation(0, 1, observational).isEmpty)
       // Otherwise, all we can assert is that -1.0 <= rho <= 1.0
-      predictionResult.getUncertaintyCorrelation(0, 2, observational).get.foreach { calcRho =>
+      predictionResult.uncertaintyCorrelation(0, 2, observational).get.foreach { calcRho =>
         assert(calcRho >= -1.0 && calcRho <= 1.0)
       }
     }
@@ -277,11 +283,10 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
   @Test
   def testSparseMixedBagged(): Unit = {
     /* Setup some data */
-    val raw: Seq[(Vector[Double], Double)] =
-      TestUtils.generateTrainingData(256, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
-    val inputs: Seq[Vector[Double]] = raw.map(_._1)
-    val realLabel: Seq[Double] = raw.map(_._2)
-    val catLabel: Seq[Boolean] = raw.map(_._2 > realLabel.max / 2.0)
+    val realRows = DataGenerator.generate(256, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng).data
+    val inputs = realRows.map(_.inputs)
+    val realLabel: Seq[Double] = realRows.map(_.label)
+    val catLabel: Seq[Boolean] = realRows.map(_.label > realLabel.max / 2.0)
     val sparseCat = catLabel.map(x =>
       if (rng.nextDouble() > 0.125) {
         null
@@ -297,6 +302,7 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
       }
     )
     val labels = Vector(sparseReal, sparseCat).transpose
+    val multiTaskRows = TrainingRow.build(inputs.zip(labels))
 
     val DTLearner = MultiTaskTreeLearner()
     val baggedLearner = MultiTaskBagger(
@@ -304,26 +310,26 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
       numBags = inputs.size,
       biasLearner = Some(GuessTheMeanLearner())
     )
-    val trainingResult = baggedLearner.train(inputs.zip(labels), rng = rng)
-    val RF = trainingResult.getModels().last
+    val trainingResult = baggedLearner.train(multiTaskRows, rng = rng)
+    val RF = trainingResult.models.last
 
-    val catResults = RF.transform(inputs).getExpected()
-    val realUncertainty = trainingResult.getModels().head.transform(inputs).getUncertainty().get
+    val catResults = RF.transform(inputs).expected
+    val realUncertainty = trainingResult.models.head.transform(inputs).uncertainty().get
     assert(realUncertainty.forall(!_.asInstanceOf[Double].isNaN), s"Some uncertainty values were NaN")
 
-    val referenceModel = Bagger(ClassificationTreeLearner(), numBags = inputs.size)
-      .train(inputs.zip(sparseCat).filterNot(_._2 == null), rng = rng)
-    val reference = referenceModel
-      .getModel()
+    val referenceRows = TrainingRow.build(inputs.zip(sparseCat)).filterNot(_.label == null)
+    val referenceModel = ClassificationBagger(ClassificationTreeLearner(), numBags = inputs.size)
+      .train(referenceRows, rng = rng)
+    val reference = referenceModel.model
       .transform(inputs)
-      .getExpected()
+      .expected
 
     val singleF1 = ClassificationMetrics.f1scores(reference, catLabel)
     val multiF1 = ClassificationMetrics.f1scores(catResults, catLabel)
 
     // Make sure we can grab the loss without issue
-    assert(!referenceModel.getLoss().get.isNaN, "Single task classification loss was NaN")
-    assert(!trainingResult.getLoss().get.isNaN, "Sparse multitask loss was NaN")
+    assert(!referenceModel.loss.get.isNaN, "Single task classification loss was NaN")
+    assert(!trainingResult.loss.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")
@@ -335,11 +341,10 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
   @Test
   def testFullSparsity(): Unit = {
     /* Setup some data */
-    val raw: Seq[(Vector[Double], Double)] =
-      TestUtils.generateTrainingData(256, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
-    val inputs: Seq[Vector[Double]] = raw.map(_._1)
-    val realLabel: Seq[Double] = raw.map(_._2)
-    val catLabel: Seq[Boolean] = raw.map(_._2 > realLabel.max / 2.0)
+    val realRows = DataGenerator.generate(256, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng).data
+    val inputs = realRows.map(_.inputs)
+    val realLabel: Seq[Double] = realRows.map(_.label)
+    val catLabel: Seq[Boolean] = realRows.map(_.label > realLabel.max / 2.0)
 
     // Test 2 real outputs, 2 categorical outputs, and an even split
     Seq(2, 128, 254).foreach { cutoffIndex =>
@@ -352,13 +357,14 @@ class MultiTaskBaggerTest extends SeedRandomMixIn {
           if (idx >= cutoffIndex) Double.NaN else x
       }
       val labels = Vector(sparseReal, sparseCat).transpose
+      val multiTaskRows = TrainingRow.build(inputs.zip(labels))
 
       val baggedLearner = MultiTaskBagger(
         MultiTaskTreeLearner(),
         numBags = inputs.size,
         biasLearner = Some(GuessTheMeanLearner())
       )
-      assert(baggedLearner.train(inputs.zip(labels), rng = rng).getModels().size == 2)
+      assert(baggedLearner.train(multiTaskRows, rng = rng).models.size == 2)
     }
   }
 }
diff --git a/src/test/scala/io/citrine/lolo/learners/ExtraRandomTreesTest.scala b/src/test/scala/io/citrine/lolo/learners/ExtraRandomTreesTest.scala
index a5112b50..cc46e011 100644
--- a/src/test/scala/io/citrine/lolo/learners/ExtraRandomTreesTest.scala
+++ b/src/test/scala/io/citrine/lolo/learners/ExtraRandomTreesTest.scala
@@ -1,11 +1,13 @@
 package io.citrine.lolo.learners
 
 import breeze.stats.distributions.{Beta, RandBasis}
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
 import io.citrine.lolo.stats.functions.Friedman
 import io.citrine.lolo.trees.regression.RegressionTreeLearner
 import io.citrine.random.Random
 import org.junit.Test
+import org.scalatest.Assertions._
 
 @Test
 class ExtraRandomTreesTest extends SeedRandomMixIn {
@@ -15,34 +17,34 @@ class ExtraRandomTreesTest extends SeedRandomMixIn {
     */
   @Test
   def testRegression(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(256, 5, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      inputBins = Seq((0, 8))
-    )
+    val trainingData = DataGenerator
+      .generate(256, 5, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .data
 
     Seq(true, false).foreach { randomlyRotateFeatures =>
-      val RFMeta = ExtraRandomTrees(randomlyRotateFeatures = randomlyRotateFeatures)
+      val RFMeta = ExtraRandomTreesRegressor(randomlyRotateFeatures = randomlyRotateFeatures)
         .train(trainingData, rng = rng)
-      val RF = RFMeta.getModel()
+      val RF = RFMeta.model
 
-      val loss = RFMeta.getLoss().get
+      val loss = RFMeta.loss.get
       assert(0 <= loss && loss < 1e-8, "Expected zero loss.")
 
-      val results = RF.transform(trainingData.map(_._1))
-      // val means = results.getExpected()
-      val sigma: Seq[Double] = results.getUncertainty(observational = false).get.asInstanceOf[Seq[Double]]
+      val results = RF.transform(trainingData.map(_.inputs))
+      // val means = results.expected
+      val sigma: Seq[Double] = results.uncertainty(observational = false).get.asInstanceOf[Seq[Double]]
       assert(sigma.forall(_ >= 0.0))
 
       assert(
-        results.getUncertainty(observational = true).isEmpty,
+        results.uncertainty(observational = true).isEmpty,
         "Observational uncertainty should be empty, since jackknife is disabled."
       )
 
-      assert(results.getGradient().isEmpty, "Returned a gradient when there shouldn't be one")
+      assert(results.gradient.isEmpty, "Returned a gradient when there shouldn't be one")
 
       if (!randomlyRotateFeatures) {
         /* The first feature should be the most important */
-        val importances = RFMeta.getFeatureImportance().get
+        val importances = RFMeta.featureImportance.get
         assert(importances(1) == importances.max)
       }
     }
@@ -56,16 +58,15 @@ class ExtraRandomTreesTest extends SeedRandomMixIn {
     val nTrain = 64
     val nTest = nTrain
     val nBins = 8
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(nTrain, 5, noise = 0.0, function = Friedman.friedmanSilverman, rng = rng),
-      responseBins = Some(nBins)
-    )
+
+    val trainingData = DataGenerator
+      .generate(nTrain, 5, noise = 0.0, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedLabels(bins = nBins)
+      .data
 
     /* Generate small perturbations from the training data */
     val testData = trainingData
-      .map { r =>
-        (r._1.map { _.asInstanceOf[Double] + 0.05 * rng.nextGaussian() }, r._2)
-      }
+      .map(_.mapInputs(_.map(x => x.asInstanceOf[Double] + 0.05 * rng.nextGaussian())))
       .take(nTest)
 
     Seq(true, false).foreach { randomlyRotateFeatures =>
@@ -73,18 +74,18 @@ class ExtraRandomTreesTest extends SeedRandomMixIn {
         Seq(1, 2).foreach { minLeafInstances =>
           rng = Random(238834L)
 
-          val RFMeta = ExtraRandomTrees(
+          val RFMeta = ExtraRandomTreesClassifier(
             numTrees = trainingData.size * 4,
             randomlyRotateFeatures = randomlyRotateFeatures,
             disableBootstrap = disableBootstrap,
             minLeafInstances = minLeafInstances
           ).train(trainingData, rng = rng)
-          val RF = RFMeta.getModel()
+          val RF = RFMeta.model
 
           /* Inspect the results on the training set */
           val trueRateTrainingSet = trainingData
-            .map(_._2)
-            .zip(RF.transform(trainingData.map(_._1)).getExpected())
+            .map(_.label)
+            .zip(RF.transform(trainingData.map(_.inputs)).expected)
             .count { case (a, p) => a == p }
             .toDouble / nTrain
           if (minLeafInstances == 1 && disableBootstrap) {
@@ -94,15 +95,15 @@ class ExtraRandomTreesTest extends SeedRandomMixIn {
           }
 
           /* Inspect the results on the perturbed test set */
-          val results = RF.transform(testData.map(_._1))
-          val means = results.getExpected()
-          val trueRateTestSet = testData.map(_._2).zip(means).count { case (a, p) => a == p }.toDouble / nTest
+          val results = RF.transform(testData.map(_.inputs))
+          val means = results.expected
+          val trueRateTestSet = testData.map(_.label).zip(means).count { case (a, p) => a == p }.toDouble / nTest
           assert(trueRateTestSet > 0.5)
 
           /* Check that class probabilities are reasonable */
-          val uncertainty = results.getUncertainty()
+          val uncertainty = results.uncertainty()
           assert(uncertainty.isDefined)
-          assert(testData.map(_._2).zip(uncertainty.get).forall {
+          assert(testData.map(_.label).zip(uncertainty.get).forall {
             case (a, probs) =>
               val classProbabilities = probs.asInstanceOf[Map[Any, Double]]
               val maxProb = classProbabilities(a)
@@ -125,34 +126,34 @@ class ExtraRandomTreesTest extends SeedRandomMixIn {
         val nTrain = 64
         val nTest = 16
         val (mainTrainingData, testData) = {
-          val allData = TestUtils.binTrainingData(
-            TestUtils.generateTrainingData(
+          val allData = DataGenerator
+            .generate(
               nTrain + nTest,
               5,
               noise = 0.1,
               function = Friedman.friedmanSilverman,
               rng = rng
-            ),
-            responseBins = Some(2)
-          )
+            )
+            .withBinnedLabels(bins = 2)
+            .data
           (allData.take(nTrain), allData.takeRight(nTest))
         }
         val dupeLabel = "DUPE"
         val trainingDataSuffixed = mainTrainingData ++ Seq(
-          (mainTrainingData.head._1, dupeLabel)
+          TrainingRow(mainTrainingData.head.inputs, dupeLabel)
         )
         val trainingDataPrefixed = Seq(
-          (mainTrainingData.head._1, dupeLabel)
+          TrainingRow(mainTrainingData.head.inputs, dupeLabel)
         ) ++ mainTrainingData
 
         val RFSuffixed =
-          ExtraRandomTrees(numTrees = trainingDataSuffixed.size * 2).train(trainingDataSuffixed, rng = rng)
+          ExtraRandomTreesClassifier(numTrees = trainingDataSuffixed.size * 2).train(trainingDataSuffixed, rng = rng)
         val RFPrefixed =
-          ExtraRandomTrees(numTrees = trainingDataPrefixed.size * 2).train(trainingDataPrefixed, rng = rng)
-        val predictedSuffixed = RFSuffixed.getModel().transform(testData.map(_._1))
-        val predictedPrefixed = RFPrefixed.getModel().transform(testData.map(_._1))
-        val extraLabelCountSuffixed = predictedSuffixed.getExpected().count { case p: String => p == dupeLabel }
-        val extraLabelCountPrefixed = predictedPrefixed.getExpected().count { case p: String => p == dupeLabel }
+          ExtraRandomTreesClassifier(numTrees = trainingDataPrefixed.size * 2).train(trainingDataPrefixed, rng = rng)
+        val predictedSuffixed = RFSuffixed.model.transform(testData.map(_.inputs))
+        val predictedPrefixed = RFPrefixed.model.transform(testData.map(_.inputs))
+        val extraLabelCountSuffixed = predictedSuffixed.expected.count { case p: String => p == dupeLabel }
+        val extraLabelCountPrefixed = predictedPrefixed.expected.count { case p: String => p == dupeLabel }
 
         if (extraLabelCountSuffixed > extraLabelCountPrefixed) {
           (1, 0)
@@ -185,12 +186,12 @@ class ExtraRandomTreesTest extends SeedRandomMixIn {
     */
   @Test
   def testWeightsWithSmallData(): Unit = {
-    val trainingData = TestUtils.generateTrainingData(8, 1)
+    val trainingData = DataGenerator.generate(8, 1).data
     // the number of trees is the number of times we generate weights
     // so this has the effect of creating lots of different sets of weights
-    val learner = ExtraRandomTrees(numTrees = 16384)
+    val learner = ExtraRandomTreesRegressor(numTrees = 16384)
     // the test is that this training doesn't throw an exception
-    learner.train(trainingData, rng = rng).getModel()
+    learner.train(trainingData, rng = rng).model
   }
 
   /** Test that the same random seed leads to identical models */
@@ -198,27 +199,30 @@ class ExtraRandomTreesTest extends SeedRandomMixIn {
   def testReproducibility(): Unit = {
     val numTrain = 50
     // Generate completely random training data
-    val (inputs: Seq[Vector[Double]], realLabel: Seq[Double]) = TestUtils
-      .generateTrainingData(rows = numTrain, cols = 12, noise = 5.0, function = _ => 0.0, rng = rng)
-      .unzip
-    val catLabel: Seq[Boolean] = Seq.fill(numTrain)(rng.nextBoolean())
+    val realRows = DataGenerator.generate(rows = numTrain, cols = 12, noise = 5.0, function = _ => 0.0, rng = rng).data
+    val catRows = realRows.map(_.withLabel(rng.nextBoolean()))
 
     // Generate test points
     val numTest = 25
-    val testInputs = TestUtils.generateTrainingData(rows = numTest, cols = 12, function = _ => 0.0, rng = rng).map(_._1)
+    val testInputs =
+      DataGenerator.generate(rows = numTest, cols = 12, function = _ => 0.0, rng = rng).data.map(_.inputs)
 
     val seed = 67852103L
-    val RFMeta = ExtraRandomTrees(
+    val extraRegressor = ExtraRandomTreesRegressor(
       biasLearner = Some(RegressionTreeLearner(maxDepth = 5)),
       randomlyRotateFeatures = true
     )
-    Vector(realLabel, catLabel).foreach { labels =>
-      val model1 = RFMeta.train(inputs.zip(labels), rng = Random(seed)).getModel()
-      val model2 = RFMeta.train(inputs.zip(labels), rng = Random(seed)).getModel()
-      val predictions1 = model1.transform(testInputs)
-      val predictions2 = model2.transform(testInputs)
-      assert(predictions1.getExpected() == predictions2.getExpected())
-    }
+    val regModel1 = extraRegressor.train(realRows, rng = Random(seed)).model
+    val regModel2 = extraRegressor.train(realRows, rng = Random(seed)).model
+    val regPredictions1 = regModel1.transform(testInputs)
+    val regPredictions2 = regModel2.transform(testInputs)
+    assert(regPredictions1.expected == regPredictions2.expected)
+
+    val extraClassifier = ExtraRandomTreesClassifier(randomlyRotateFeatures = true)
+    val classModel1 = extraClassifier.train(catRows, rng = Random(seed)).model
+    val classModel2 = extraClassifier.train(catRows, rng = Random(seed)).model
+    val classPredictions1 = classModel1.transform(testInputs)
+    val classPredictions2 = classModel2.transform(testInputs)
+    assert(classPredictions1.expected == classPredictions2.expected)
   }
-
 }
diff --git a/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala b/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala
index 1566a9d3..91c99c7e 100644
--- a/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala
+++ b/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala
@@ -2,12 +2,11 @@ package io.citrine.lolo.learners
 
 import breeze.linalg.DenseMatrix
 import breeze.stats.distributions.{Beta, RandBasis}
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
-import io.citrine.lolo.bags.MultiTaskBaggedResult
-import io.citrine.lolo.linear.GuessTheMeanLearner
+import io.citrine.lolo.api.{Learner, TrainingRow}
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
+import io.citrine.lolo.bags.MultiTaskBaggedPrediction
 import io.citrine.lolo.stats.functions.Friedman
 import io.citrine.lolo.trees.regression.RegressionTreeLearner
-import io.citrine.lolo.trees.splits.RegressionSplitter
 import io.citrine.random.Random
 import org.junit.Test
 import org.scalatest.Assertions._
@@ -20,28 +19,28 @@ class RandomForestTest extends SeedRandomMixIn {
     */
   @Test
   def testRegressionForest(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      inputBins = Seq((0, 8))
-    )
+    val trainingData = DataGenerator
+      .generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .data
 
     Seq(true, false).foreach { randomlyRotateFeatures =>
-      val RFMeta = RandomForest(randomlyRotateFeatures = randomlyRotateFeatures)
+      val RFMeta = RandomForestRegressor(randomlyRotateFeatures = randomlyRotateFeatures)
         .train(trainingData)
-      val RF = RFMeta.getModel()
+      val RF = RFMeta.model
 
-      assert(RFMeta.getLoss().get < 1.0, "Loss of bagger is larger than expected")
+      assert(RFMeta.loss.get < 1.0, "Loss of bagger is larger than expected")
 
-      val results = RF.transform(trainingData.map(_._1))
-      // val means = results.getExpected()
-      val sigma: Seq[Double] = results.getUncertainty().get.asInstanceOf[Seq[Double]]
+      val results = RF.transform(trainingData.map(_.inputs))
+      // val means = results.expected
+      val sigma: Seq[Double] = results.uncertainty().get.asInstanceOf[Seq[Double]]
       assert(sigma.forall(_ >= 0.0))
 
-      assert(results.getGradient().isEmpty, "Returned a gradient when there shouldn't be one")
+      assert(results.gradient.isEmpty, "Returned a gradient when there shouldn't be one")
 
       if (!randomlyRotateFeatures) {
         /* The first feature should be the most important */
-        val importances = RFMeta.getFeatureImportance().get
+        val importances = RFMeta.featureImportance.get
         assert(importances(1) == importances.max)
       }
     }
@@ -49,24 +48,25 @@ class RandomForestTest extends SeedRandomMixIn {
 
   /** 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, rng = rng),
-        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().train(inputs.zip(labels), rng = rng)
-    val model = RFMeta.getModel()
-
-    val results = model.transform(inputs).asInstanceOf[MultiTaskBaggedResult]
-    assert(results.getUncertainty().isDefined)
-    assert(results.getUncertaintyCorrelation(0, 2).isDefined)
+  def testMultiTaskForest(): Unit = {
+    val realRows = DataGenerator
+      .generate(256, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .data
+
+    val inputs = realRows.map(_.inputs)
+    val realLabels = realRows.map(_.label)
+    val catLabels = realLabels.map(_ > realLabels.max / 2.0)
+    val quadLabels = realLabels.map(x => x * x)
+    val allLabels = Vector(realLabels, catLabels, quadLabels).transpose
+    val multiTaskRows = TrainingRow.build(inputs.zip(allLabels))
+
+    val RFMeta = MultiTaskRandomForest().train(multiTaskRows, rng = rng)
+    val model = RFMeta.model
+
+    val results = model.transform(inputs).asInstanceOf[MultiTaskBaggedPrediction]
+    assert(results.uncertainty().isDefined)
+    assert(results.uncertaintyCorrelation(0, 2).isDefined)
   }
 
   /**
@@ -74,27 +74,26 @@ class RandomForestTest extends SeedRandomMixIn {
     */
   @Test
   def testClassificationForest(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils
-        .generateTrainingData(128, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      inputBins = Seq((0, 8)),
-      responseBins = Some(8)
-    )
+    val trainingData = DataGenerator
+      .generate(128, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .withBinnedLabels(bins = 8)
+      .data
 
     Seq(true, false).foreach { randomlyRotateFeatures =>
       val RFMeta =
-        RandomForest(numTrees = trainingData.size * 2, randomlyRotateFeatures = randomlyRotateFeatures)
+        RandomForestClassifier(numTrees = trainingData.size * 2, randomlyRotateFeatures = randomlyRotateFeatures)
           .train(trainingData, rng = rng)
-      val RF = RFMeta.getModel()
+      val RF = RFMeta.model
 
       /* Inspect the results */
-      val results = RF.transform(trainingData.map(_._1))
-      val means = results.getExpected()
-      assert(trainingData.map(_._2).zip(means).forall { case (a, p) => a == p })
+      val results = RF.transform(trainingData.map(_.inputs))
+      val means = results.expected
+      assert(trainingData.map(_.label).zip(means).forall { case (a, p) => a == p })
 
-      val uncertainty = results.getUncertainty()
+      val uncertainty = results.uncertainty()
       assert(uncertainty.isDefined)
-      assert(trainingData.map(_._2).zip(uncertainty.get).forall {
+      assert(trainingData.map(_.label).zip(uncertainty.get).forall {
         case (a, probs) =>
           val classProbabilities = probs.asInstanceOf[Map[Any, Double]]
           val maxProb = classProbabilities(a)
@@ -106,31 +105,53 @@ class RandomForestTest extends SeedRandomMixIn {
   /** Test that the same random seed leads to identical models */
   @Test
   def testReproducibility(): Unit = {
+    def checkReproducibility[T](
+        learner: Learner[T],
+        trainingData: Seq[TrainingRow[T]],
+        testInputs: Seq[Vector[Any]],
+        seed: Long
+    ): Unit = {
+      val model1 = learner.train(trainingData, rng = Random(seed)).model
+      val model2 = learner.train(trainingData, rng = Random(seed)).model
+      val predictions1 = model1.transform(testInputs)
+      val predictions2 = model2.transform(testInputs)
+      assert(predictions1.expected == predictions2.expected)
+    }
+
     val numTrain = 50
     // Generate completely random training data
-    val (inputs: Seq[Vector[Double]], realLabel: Seq[Double]) = TestUtils
-      .generateTrainingData(rows = numTrain, cols = 12, noise = 5.0, function = _ => 0.0, rng = rng)
-      .unzip
-    val catLabel: Seq[Boolean] = Seq.fill(numTrain)(rng.nextBoolean())
-    val allLabels = Vector(realLabel, catLabel).transpose
+    val realRows = DataGenerator
+      .generate(rows = numTrain, cols = 12, noise = 5.0, function = _ => 0.0, rng = rng)
+      .data
+    val catRows = realRows.map(_.withLabel(rng.nextBoolean()))
+    val multiTaskRows = realRows.zip(catRows).map {
+      case (r, c) => r.withLabel(Vector(r.label, c.label))
+    }
 
     // Generate test points
     val numTest = 25
-    val testInputs = TestUtils.generateTrainingData(rows = numTest, cols = 12, function = _ => 0.0, rng = rng).map(_._1)
+    val testInputs =
+      DataGenerator.generate(rows = numTest, cols = 12, function = _ => 0.0, rng = rng).data.map(_.inputs)
 
     val seed = 67852103L
-    val RFMeta = RandomForest(
+    val rfRegressor = RandomForestRegressor(
       biasLearner = Some(RegressionTreeLearner(maxDepth = 5)),
       randomizePivotLocation = true,
       randomlyRotateFeatures = true
     )
-    Vector(realLabel, catLabel, allLabels).foreach { labels =>
-      val model1 = RFMeta.train(inputs.zip(labels), rng = Random(seed)).getModel()
-      val model2 = RFMeta.train(inputs.zip(labels), rng = Random(seed)).getModel()
-      val predictions1 = model1.transform(testInputs)
-      val predictions2 = model2.transform(testInputs)
-      assert(predictions1.getExpected() == predictions2.getExpected())
-    }
+    checkReproducibility(rfRegressor, realRows, testInputs, seed)
+
+    val rfClassifier = RandomForestClassifier(
+      randomizePivotLocation = true,
+      randomlyRotateFeatures = true
+    )
+    checkReproducibility(rfClassifier, catRows, testInputs, seed)
+
+    val rfMultiTask = MultiTaskRandomForest(
+      randomizePivotLocation = true,
+      randomlyRotateFeatures = true
+    )
+    checkReproducibility(rfMultiTask, multiTaskRows, testInputs, seed)
   }
 
   /**
@@ -142,25 +163,27 @@ class RandomForestTest extends SeedRandomMixIn {
     val numTrials = 20
     val (winsSuffixed, winsPrefixed): (Int, Int) = (0 until numTrials)
       .map { _ =>
-        val mainTrainingData = TestUtils.binTrainingData(
-          TestUtils
-            .generateTrainingData(64, 5, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-          responseBins = Some(2)
-        )
+        val mainTrainingData = DataGenerator
+          .generate(64, 5, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+          .withBinnedLabels(bins = 2)
+          .data
+
         val dupeLabel = "DUPE"
         val trainingDataSuffixed = mainTrainingData ++ Seq(
-          (mainTrainingData.head._1, dupeLabel)
+          TrainingRow(mainTrainingData.head.inputs, dupeLabel)
         )
         val trainingDataPrefixed = Seq(
-          (mainTrainingData.head._1, dupeLabel)
+          TrainingRow(mainTrainingData.head.inputs, dupeLabel)
         ) ++ mainTrainingData
 
-        val RFSuffixed = RandomForest(numTrees = trainingDataSuffixed.size * 2).train(trainingDataSuffixed, rng = rng)
-        val RFPrefixed = RandomForest(numTrees = trainingDataPrefixed.size * 2).train(trainingDataPrefixed, rng = rng)
-        val predictedSuffixed = RFSuffixed.getModel().transform(mainTrainingData.map(_._1))
-        val predictedPrefixed = RFPrefixed.getModel().transform(mainTrainingData.map(_._1))
-        val extraLabelCountSuffixed = predictedSuffixed.getExpected().count { case p: String => p == dupeLabel }
-        val extraLabelCountPrefixed = predictedPrefixed.getExpected().count { case p: String => p == dupeLabel }
+        val RFSuffixed =
+          RandomForestClassifier(numTrees = trainingDataSuffixed.size * 2).train(trainingDataSuffixed, rng = rng)
+        val RFPrefixed =
+          RandomForestClassifier(numTrees = trainingDataPrefixed.size * 2).train(trainingDataPrefixed, rng = rng)
+        val predictedSuffixed = RFSuffixed.model.transform(mainTrainingData.map(_.inputs))
+        val predictedPrefixed = RFPrefixed.model.transform(mainTrainingData.map(_.inputs))
+        val extraLabelCountSuffixed = predictedSuffixed.expected.count { case p: String => p == dupeLabel }
+        val extraLabelCountPrefixed = predictedPrefixed.expected.count { case p: String => p == dupeLabel }
 
         if (extraLabelCountSuffixed > extraLabelCountPrefixed) {
           (1, 0)
@@ -195,28 +218,25 @@ class RandomForestTest extends SeedRandomMixIn {
   @Test
   def testRandomizedSplitLocations(): Unit = {
     // Generate a linear signal in one dimension: 2 * x
-    val trainingData: Seq[(Vector[Double], Double)] = TestUtils.generateTrainingData(
-      32,
-      1,
-      function = { x => x.head * 2.0 },
-      rng = rng
-    )
+    val trainingData = DataGenerator
+      .generate(32, 1, function = { x => x.head * 2.0 }, rng = rng)
+      .data
 
     // Create a consistent set of parameters
-    val baseForest = RandomForest(numTrees = 16384, useJackknife = false)
+    val baseForest = RandomForestRegressor(numTrees = 16384, useJackknife = false)
 
     // Turn off split randomization and compute the loss (out-of-bag error)
     val lossWithoutRandomization: Double = baseForest
       .copy(randomizePivotLocation = false)
       .train(trainingData, rng = rng)
-      .getLoss()
+      .loss
       .get
 
     // Turn on split randomization and compute the loss (out-of-bag error)
     val lossWithRandomization: Double = baseForest
       .copy(randomizePivotLocation = true)
       .train(trainingData, rng = rng)
-      .getLoss()
+      .loss
       .get
 
     assert(lossWithRandomization < lossWithoutRandomization)
@@ -227,25 +247,25 @@ class RandomForestTest extends SeedRandomMixIn {
     */
   @Test
   def testWeightsWithSmallData(): Unit = {
-    val trainingData = TestUtils.generateTrainingData(8, 1, rng = rng)
+    val trainingData = DataGenerator.generate(8, 1, rng = rng).data
     // the number of trees is the number of times we generate weights
     // so this has the effect of creating lots of different sets of weights
-    val learner = RandomForest(numTrees = 16384)
+    val learner = RandomForestRegressor(numTrees = 16384)
     // the test is that this training doesn't throw an exception
-    learner.train(trainingData, rng = rng).getModel()
+    learner.train(trainingData, rng = rng).model
   }
 
   def shapleyCompare(
-      trainingData: Seq[(Vector[Any], Double)],
+      trainingData: Seq[TrainingRow[Double]],
       evalLocation: Vector[Any],
       expected: Vector[Double],
       rtol: Double = 0.1
   ): Unit = {
-    val actual = RandomForest().train(trainingData, rng = rng).getModel().shapley(evalLocation) match {
+    val actual = RandomForestRegressor().train(trainingData, rng = rng).model.shapley(evalLocation) match {
       case None => fail("Unexpected None returned by shapley.")
       case x: Option[DenseMatrix[Double]] => {
         val a = x.get
-        assert(a.cols == trainingData.head._1.length, "Expected one Shapley value per feature.")
+        assert(a.cols == trainingData.head.inputs.length, "Expected one Shapley value per feature.")
         assert(a.rows == 1, "Expected a single output dimension.")
         a.toDenseVector.toScalaVector
       }
@@ -265,48 +285,48 @@ class RandomForestTest extends SeedRandomMixIn {
 
     // Example from Lundberg paper (https://arxiv.org/pdf/1802.03888.pdf)
     val trainingData1 = Seq(
-      (Vector(1.0, 1.0), 80.0),
-      (Vector(1.0, 0.0), 0.0),
-      (Vector(0.0, 1.0), 0.0),
-      (Vector(0.0, 0.0), 0.0)
+      TrainingRow(Vector(1.0, 1.0), 80.0),
+      TrainingRow(Vector(1.0, 0.0), 0.0),
+      TrainingRow(Vector(0.0, 1.0), 0.0),
+      TrainingRow(Vector(0.0, 0.0), 0.0)
     )
     val expected1 = Vector(30.0, 30.0)
-    shapleyCompare((1 to 8).map(_ => trainingData1).flatten, Vector[Any](1.0, 1.0), expected1)
+    shapleyCompare((1 to 8).flatMap(_ => trainingData1), Vector[Any](1.0, 1.0), expected1)
 
     // Second example from Lundberg paper (https://arxiv.org/pdf/1802.03888.pdf)
     val trainingData2 = Seq(
-      (Vector(1.0, 1.0), 90.0),
-      (Vector(1.0, 0.0), 10.0),
-      (Vector(0.0, 1.0), 0.0),
-      (Vector(0.0, 0.0), 0.0)
+      TrainingRow(Vector(1.0, 1.0), 90.0),
+      TrainingRow(Vector(1.0, 0.0), 10.0),
+      TrainingRow(Vector(0.0, 1.0), 0.0),
+      TrainingRow(Vector(0.0, 0.0), 0.0)
     )
     val expected2 = Vector(35.0, 30.0)
-    shapleyCompare((1 to 8).map(_ => trainingData2).flatten, Vector[Any](1.0, 1.0), expected2)
+    shapleyCompare((1 to 8).flatMap(_ => trainingData2), Vector[Any](1.0, 1.0), expected2)
 
     // Example with two splits on one feature
     // Worked out with pen-and-paper from Lundberg Equation 2.
     val trainingData3 = Seq(
-      (Vector(1.0, 1.0), 100.0),
-      (Vector(1.0, 0.0), 80.0),
-      (Vector(1.0, 0.2), 70.0),
-      (Vector(0.0, 1.0), 0.0),
-      (Vector(0.0, 0.2), 0.0),
-      (Vector(0.0, 0.0), 0.0)
+      TrainingRow(Vector(1.0, 1.0), 100.0),
+      TrainingRow(Vector(1.0, 0.0), 80.0),
+      TrainingRow(Vector(1.0, 0.2), 70.0),
+      TrainingRow(Vector(0.0, 1.0), 0.0),
+      TrainingRow(Vector(0.0, 0.2), 0.0),
+      TrainingRow(Vector(0.0, 0.0), 0.0)
     )
     val expected3 = Vector(45.8333333333333, 12.5)
-    shapleyCompare((1 to 8).map(_ => trainingData3).flatten, Vector[Any](1.0, 1.0), expected3)
+    shapleyCompare((1 to 8).flatMap(_ => trainingData3), Vector[Any](1.0, 1.0), expected3)
 
     // Example with 5 features, to exercise all the factorials in Lundberg equation 2.
     // Referenced against the shap package on a sklearn decision tree.
     val trainingData4 = Seq(
-      (Vector(0.0, 0.0, 0.0, 0.0, 0.0), 1.0),
-      (Vector(1.0, 0.0, 0.0, 0.0, 0.0), 2.0),
-      (Vector(0.0, 1.0, 0.0, 0.0, 0.0), 4.0),
-      (Vector(0.0, 0.0, 1.0, 0.0, 0.0), 8.0),
-      (Vector(0.0, 0.0, 0.0, 1.0, 0.0), 16.0),
-      (Vector(0.0, 0.0, 0.0, 0.0, 1.0), 32.0)
+      TrainingRow(Vector(0.0, 0.0, 0.0, 0.0, 0.0), 1.0),
+      TrainingRow(Vector(1.0, 0.0, 0.0, 0.0, 0.0), 2.0),
+      TrainingRow(Vector(0.0, 1.0, 0.0, 0.0, 0.0), 4.0),
+      TrainingRow(Vector(0.0, 0.0, 1.0, 0.0, 0.0), 8.0),
+      TrainingRow(Vector(0.0, 0.0, 0.0, 1.0, 0.0), 16.0),
+      TrainingRow(Vector(0.0, 0.0, 0.0, 0.0, 1.0), 32.0)
     )
     val expected4 = Vector(0.0333333333333333, 0.2, 0.8666666666666667, 3.533333333333333, 16.866666666666667)
-    shapleyCompare((1 to 8).map { _ => trainingData4 }.flatten, Vector.fill[Any](5)(1.0), expected4)
+    shapleyCompare((1 to 8).flatMap { _ => trainingData4 }, Vector.fill[Any](5)(1.0), expected4)
   }
 }
diff --git a/src/test/scala/io/citrine/lolo/linear/LinearRegressionTest.scala b/src/test/scala/io/citrine/lolo/linear/LinearRegressionTest.scala
index 3ba8e318..f8b2aa07 100644
--- a/src/test/scala/io/citrine/lolo/linear/LinearRegressionTest.scala
+++ b/src/test/scala/io/citrine/lolo/linear/LinearRegressionTest.scala
@@ -2,6 +2,7 @@ package io.citrine.lolo.linear
 
 import breeze.linalg.{norm, DenseMatrix, DenseVector}
 import io.citrine.lolo.SeedRandomMixIn
+import io.citrine.lolo.api.TrainingRow
 import org.junit.Test
 
 /**
@@ -30,15 +31,15 @@ class LinearRegressionTest extends SeedRandomMixIn {
     val result = data * beta0
 
     val trainingData = (0 until n).map { i =>
-      (data.t(::, i).toDenseVector.toArray.toVector, result(i))
+      TrainingRow(data.t(::, i).toDenseVector.toArray.toVector, result(i))
     }
 
     val lr = LinearRegressionLearner(fitIntercept = false)
     val lrm = lr.train(trainingData)
-    val model = lrm.getModel()
-    val output = model.transform(trainingData.map(_._1))
-    val predicted = output.getExpected()
-    val beta = output.getGradient().get.head
+    val model = lrm.model
+    val output = model.transform(trainingData.map(_.inputs))
+    val predicted = output.expected
+    val beta = output.gradient.get.head
 
     assert(norm(new DenseVector(beta.toArray) - beta0) < 1.0e-9, "Coefficients are inaccurate")
     assert(norm(new DenseVector(predicted.toArray) - result) < 1.0e-9, "Predictions are inaccurate")
@@ -53,15 +54,15 @@ class LinearRegressionTest extends SeedRandomMixIn {
     val result = data * beta0 + int0
 
     val trainingData = (0 until n).map { i =>
-      (data.t(::, i).toDenseVector.toArray.toVector, result(i))
+      TrainingRow(data.t(::, i).toDenseVector.toArray.toVector, result(i))
     }
 
     val lr = LinearRegressionLearner()
     val lrm = lr.train(trainingData)
-    val model = lrm.getModel()
-    val output = model.transform(trainingData.map(_._1))
-    val predicted = output.getExpected()
-    val beta = output.getGradient().get.head
+    val model = lrm.model
+    val output = model.transform(trainingData.map(_.inputs))
+    val predicted = output.expected
+    val beta = output.gradient.get.head
 
     assert(norm(new DenseVector(beta.toArray) - beta0) < 1.0e-9, "Coefficients are inaccurate")
     assert(norm(new DenseVector(predicted.toArray) - result) < 1.0e-9, "Predictions are inaccurate")
@@ -76,15 +77,15 @@ class LinearRegressionTest extends SeedRandomMixIn {
     val result = data * beta0 + int0
 
     val trainingData = (0 until n).map { i =>
-      (data.t(::, i).toDenseVector.toArray.toVector, result(i))
+      TrainingRow(data.t(::, i).toDenseVector.toArray.toVector, result(i), math.abs(rng.nextDouble()))
     }
 
     val lr = LinearRegressionLearner()
-    val lrm = lr.train(trainingData, weights = Some(Seq.tabulate(n)(i => Math.abs(rng.nextDouble()))))
-    val model = lrm.getModel()
-    val output = model.transform(trainingData.map(_._1))
-    val predicted = output.getExpected()
-    val beta = output.getGradient().get.head
+    val lrm = lr.train(trainingData, rng = rng)
+    val model = lrm.model
+    val output = model.transform(trainingData.map(_.inputs))
+    val predicted = output.expected
+    val beta = output.gradient.get.head
 
     assert(norm(new DenseVector(beta.toArray) - beta0) < 1.0e-9, "Coefficients are inaccurate")
     assert(norm(new DenseVector(predicted.toArray) - result) < 1.0e-9, "Predictions are inaccurate")
@@ -110,15 +111,15 @@ class LinearRegressionTest extends SeedRandomMixIn {
     val result = data * beta0
 
     val trainingData = (0 until n).map { i =>
-      (data.t(::, i).toDenseVector.toArray.toVector, result(i))
+      TrainingRow(data.t(::, i).toDenseVector.toArray.toVector, result(i))
     }
 
-    val lr = new LinearRegressionLearner(fitIntercept = false)
+    val lr = LinearRegressionLearner(fitIntercept = false)
     val lrm = lr.train(trainingData)
-    val model = lrm.getModel()
-    val output = model.transform(trainingData.map(_._1))
-    val predicted = output.getExpected()
-    val beta = output.getGradient().head
+    val model = lrm.model
+    val output = model.transform(trainingData.map(_.inputs))
+    val predicted = output.expected
+    val beta = output.gradient.head
 
     assert(norm(new DenseVector(predicted.toArray) - result) < 1.0e-9, "Predictions are inaccurate")
   }
@@ -132,16 +133,17 @@ class LinearRegressionTest extends SeedRandomMixIn {
     val result = data * beta0 + int0
 
     val trainingData = (0 until n).map { i =>
-      ("Foo" +: data.t(::, i).toDenseVector.toArray.toVector.asInstanceOf[Vector[Any]] :+ "Bar", result(i))
+      val inputs = "Foo" +: data.t(::, i).toDenseVector.toArray.toVector.asInstanceOf[Vector[Any]] :+ "Bar"
+      TrainingRow(inputs, result(i))
     }
 
-    val lr = new LinearRegressionLearner()
+    val lr = LinearRegressionLearner()
     val lrm = lr.train(trainingData)
-    val model = lrm.getModel()
-    val output = model.transform(trainingData.map(_._1))
-    val predicted = output.getExpected()
-    val beta = output.getGradient().get.head
-    val importance = lrm.getFeatureImportance().get
+    val model = lrm.model
+    val output = model.transform(trainingData.map(_.inputs))
+    val predicted = output.expected
+    val beta = output.gradient.get.head
+    val importance = lrm.featureImportance.get
 
     /* Make sure that feature importance matches the gradient */
     val betaScale = beta.map(Math.abs).sum
diff --git a/src/test/scala/io/citrine/lolo/transformers/FeatureRotatorTest.scala b/src/test/scala/io/citrine/lolo/transformers/FeatureRotatorTest.scala
deleted file mode 100644
index a249be20..00000000
--- a/src/test/scala/io/citrine/lolo/transformers/FeatureRotatorTest.scala
+++ /dev/null
@@ -1,347 +0,0 @@
-package io.citrine.lolo.transformers
-
-import breeze.linalg.{det, DenseMatrix}
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
-import io.citrine.lolo.linear.{GuessTheMeanLearner, LinearRegressionLearner}
-import io.citrine.lolo.stats.functions.Friedman
-import io.citrine.lolo.stats.metrics.ClassificationMetrics
-import io.citrine.lolo.trees.classification.ClassificationTreeLearner
-import io.citrine.lolo.trees.multitask.MultiTaskTreeLearner
-import io.citrine.lolo.trees.regression.RegressionTreeLearner
-import org.junit.Test
-
-/**
-  * Created by gregor-robinson on 2020-01-06.
-  */
-@Test
-class FeatureRotatorTest extends SeedRandomMixIn {
-
-  val data: Seq[(Vector[Any], Any)] = TestUtils.binTrainingData(
-    TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-    inputBins = Seq((0, 8))
-  )
-  val weights: Vector[Double] = Vector.fill(data.size)(if (rng.nextBoolean()) rng.nextDouble() else 0.0)
-
-  @Test
-  def testRandomRotation(): Unit = {
-    val inputs = data.map(_._1)
-    val numRealInputs = inputs.head.length - 1
-    for (i <- 1 to 10) {
-      val U = FeatureRotator.getRandomRotation(numRealInputs)
-      assert(U.rows == numRealInputs)
-      assert(U.cols == numRealInputs)
-
-      // Check that the matrix is unitary.
-      (U.t * U - DenseMatrix.eye[Double](numRealInputs)).toArray.foreach { x =>
-        assert(Math.abs(x) < 1e-9)
-      }
-
-      // Check that the determinant is one.
-      val d = det(U)
-      assert(Math.abs(d - 1.0) < 1e-9, s"Determinant of U should be 1 but is ${d}")
-    }
-  }
-
-  @Test
-  def testApplyRotation(): Unit = {
-    val inputs = data.map(_._1)
-    val numRealInputs = inputs.head.length - 1
-    val featuresToRotate = (1 to numRealInputs).asInstanceOf[IndexedSeq[Int]]
-
-    // Create rotation matrix that exchanges first, second, and last real inputs
-    val U = DenseMatrix.eye[Double](dim = numRealInputs)
-    U(0, 0) = 0.0
-    U(1, 1) = 0.0
-    U(numRealInputs - 1, numRealInputs - 1) = 0.0
-    U(0, 1) = 1.0
-    U(1, numRealInputs - 1) = 1.0
-    U(numRealInputs - 1, 0) = 1.0
-
-    val rotatedInputs = FeatureRotator.applyRotation(inputs, featuresToRotate, U)
-    inputs.indices.foreach { i =>
-      assert(inputs(i)(0) == rotatedInputs(i)(0), "Failed to leave categorical invariant.")
-      assert(
-        Math.abs(inputs(i)(1).asInstanceOf[Double] - rotatedInputs(i)(numRealInputs).asInstanceOf[Double]) < 1e-9,
-        "Failed to exchange coordinates."
-      )
-      assert(
-        Math.abs(inputs(i)(2).asInstanceOf[Double] - rotatedInputs(i)(1).asInstanceOf[Double]) < 1e-9,
-        "Failed to exchange coordinates."
-      )
-      assert(
-        Math.abs(inputs(i)(numRealInputs).asInstanceOf[Double] - rotatedInputs(i)(2).asInstanceOf[Double]) < 1e-9,
-        "Failed to exchange coordinates."
-      )
-      (3 to numRealInputs - 1).foreach { j =>
-        assert(
-          Math.abs(inputs(i)(j).asInstanceOf[Double] - rotatedInputs(i)(j).asInstanceOf[Double]) < 1e-9,
-          "Unexpected modification of un-rotated coordinate."
-        )
-      }
-    }
-
-    // Check that we can undo rotation by applying the transpose
-    val unrotatedInputs = FeatureRotator.applyRotation(rotatedInputs, featuresToRotate, U.t)
-    inputs.indices.foreach { i =>
-      assert(inputs(i)(0) == unrotatedInputs(i)(0), "Failed to leave categorical invariant.")
-      (1 to numRealInputs - 1).foreach { j =>
-        assert(
-          Math.abs(inputs(i)(j).asInstanceOf[Double] - unrotatedInputs(i)(j).asInstanceOf[Double]) < 1e-9,
-          "Unexpected modification of coordinate after applying transpose."
-        )
-      }
-    }
-  }
-
-  /**
-    * Functions getLoss and getPredictedVsActual should pass through to base learner.
-    */
-  @Test
-  def testPassthroughFunctions(): Unit = {
-    val rotatedTrainingResult = FeatureRotator(RegressionTreeLearner()).train(data, rng = rng)
-
-    assert(
-      rotatedTrainingResult.getLoss() == rotatedTrainingResult.baseTrainingResult.getLoss(),
-      "Function getLoss() should pass through to base learner."
-    )
-
-    rotatedTrainingResult.getPredictedVsActual().foreach { x =>
-      x.zip(data).foreach {
-        case (a, b) =>
-          assert(a._1 == b._1, "getPredictedVsActual must return the correct training inputs.")
-          assert(a._2 == b._2, "getPredictedVsActual must return the correct predicted value.")
-          assert(a._3 == b._2, "getPredictedVsActual must return the correct actual value.")
-      }
-    }
-  }
-
-  /**
-    * Guess the mean should be invariant under rotation.
-    */
-  @Test
-  def testRotatedGTM(): Unit = {
-    val learner = GuessTheMeanLearner()
-    val model = learner.train(data, rng = rng).getModel()
-    val result = model.transform(data.map(_._1)).getExpected()
-
-    val rotatedLearner = FeatureRotator(GuessTheMeanLearner())
-    val rotatedModel = rotatedLearner.train(data, rng = rng).getModel()
-    val rotatedResult = rotatedModel.transform(data.map(_._1)).getExpected()
-
-    result.zip(rotatedResult).foreach {
-      case (free: Double, rotated: Double) =>
-        assert(Math.abs(free - rotated) < 1.0e-9, s"${free} and ${rotated} should be the same")
-    }
-  }
-
-  /**
-    * Linear regression should be invariant under rotation
-    */
-  @Test
-  def testRotatedLinear(): Unit = {
-    val learner = LinearRegressionLearner()
-    val model = learner.train(data, Some(weights), rng).getModel()
-    val result = model.transform(data.map(_._1))
-    val expected = result.getExpected()
-    val gradient = result.getGradient()
-
-    val rotatedLearner = FeatureRotator(learner)
-    val rotatedModel = rotatedLearner.train(data, Some(weights), rng).getModel()
-    val rotatedResult = rotatedModel.transform(data.map(_._1))
-    val rotatedExpected = rotatedResult.getExpected()
-    val rotatedGradient = rotatedResult.getGradient()
-
-    expected.zip(rotatedExpected).foreach {
-      case (free: Double, rotated: Double) =>
-        assert(Math.abs(free - rotated) < 1.0e-9, s"${free} and ${rotated} should be the same")
-    }
-
-    gradient.get.zip(rotatedGradient.get).foreach {
-      case (free, rotated) =>
-        val diff = free.zip(rotated).map { case (f, s) => Math.abs(f - s) }.max
-        assert(diff < 1.0e-9, "Gradients should be the same")
-    }
-  }
-
-  /**
-    * Ridge regression should not depend on rotation
-    */
-  @Test
-  def testRotatedRidge(): Unit = {
-    val learner = LinearRegressionLearner(regParam = Some(1.0))
-    val model = learner.train(data, rng = rng).getModel()
-    val result = model.transform(data.map(_._1)).getExpected()
-
-    val rotatedLearner = new FeatureRotator(learner)
-    val rotatedModel = rotatedLearner.train(data, rng = rng).getModel()
-    val rotatedResult = rotatedModel.transform(data.map(_._1)).getExpected()
-
-    result.zip(rotatedResult).foreach {
-      case (free: Double, rotated: Double) =>
-        assert(Math.abs(free - rotated) < 1.0e-9, s"${free} and ${rotated} should be the same")
-    }
-  }
-
-  /**
-    * Verify expected properties of rotated-input regression trees
-    */
-  @Test
-  def testRotatedRegressionTree(): Unit = {
-    val learner = RegressionTreeLearner()
-    val model = learner.train(data, rng = rng).getModel()
-    val result = model.transform(data.map(_._1)).getExpected()
-
-    val rotatedLearner = new FeatureRotator(learner)
-    val rotatedModel = rotatedLearner.train(data, rng = rng).getModel().asInstanceOf[RotatedFeatureModel[Double]]
-    var rotatedResult = rotatedModel.transform(data.map(_._1)).getExpected()
-    result.zip(rotatedResult).foreach {
-      case (free: Double, rotated: Double) =>
-        assert(Math.abs(free - rotated) < 1.0e-9, s"${free} and ${rotated} should be the same")
-    }
-
-    val rotatedData = FeatureRotator.applyRotation(data.map(_._1), rotatedModel.rotatedFeatures, rotatedModel.trans)
-    rotatedResult = rotatedModel.transform(rotatedData).getExpected()
-    // Check that labels change when we feed in different data.
-    assert(
-      rotatedResult.zip(result).map { case (a: Double, b: Double) => a - b }.count { x => Math.abs(x) > 1e-9 } > 0,
-      "Rotated data passed to rotated model should not map to the same predictions."
-    )
-    val baseResult = rotatedModel.baseModel.transform(rotatedData).getExpected()
-    // Check that labels are the same as feeding rotated data into base learner.
-    assert(
-      baseResult.zip(result).map { case (a: Double, b: Double) => a - b }.count { x => Math.abs(x) > 1e-9 } == 0,
-      "Rotated data passed to base model should map to the same predictions."
-    )
-  }
-
-  /**
-    * Verify expected properties of rotated-input classification trees
-    */
-  @Test
-  def testRotatedClassificationTree(): Unit = {
-    val classificationData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(2048, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      responseBins = Some(2)
-    )
-
-    val learner = ClassificationTreeLearner()
-    val model = learner.train(classificationData, rng = rng).getModel()
-    val result = model.transform(classificationData.map(_._1)).getExpected()
-
-    val rotatedLearner = new FeatureRotator(learner)
-    val rotatedModel =
-      rotatedLearner.train(classificationData, rng = rng).getModel().asInstanceOf[RotatedFeatureModel[String]]
-    var rotatedResult = rotatedModel.transform(classificationData.map(_._1)).getExpected()
-
-    result.zip(rotatedResult).foreach {
-      case (free: Any, rotated: Any) =>
-        assert(free == rotated, s"${free} and ${rotated} should be the same")
-    }
-
-    val rotatedData =
-      FeatureRotator.applyRotation(classificationData.map(_._1), rotatedModel.rotatedFeatures, rotatedModel.trans)
-    rotatedResult = rotatedModel.transform(rotatedData).getExpected()
-    // Check that labels change when we feed in different data.
-    assert(
-      rotatedResult.zip(result).count { case (a: String, b: String) => a != b } > 0,
-      "Rotated data passed to rotated model should not map to the same predictions."
-    )
-    val baseResult = rotatedModel.baseModel.transform(rotatedData).getExpected()
-    // Check that labels are the same as feeding rotated data into base learner.
-    assert(
-      baseResult.zip(result).count { case (a: String, b: String) => a != b } == 0,
-      s"Categorical labels should be identical when we feed rotated inputs into base learner."
-    )
-  }
-
-  /**
-    * Verify that rotated-input multi-task trees have expected properties
-    */
-  @Test
-  def testMultiTaskRotator(): Unit = {
-    val data: Vector[(Vector[Double], Double)] =
-      TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
-
-    // Generate multi-task training data
-    val (inputs, doubleLabel) = data.unzip
-    val catLabel = inputs.map(x => Friedman.friedmanGrosseSilverman(x) > 15.0)
-
-    // 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 = MultiTaskTreeLearner()
-    val rotatedLearner = MultiTaskFeatureRotator(MultiTaskTreeLearner())
-
-    val baseTrainingResult = baseLearner.train(inputs.zip(labels), rng = rng)
-    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()
-    val rotatedDoubleResult = rotatedDoubleModel.transform(inputs).getExpected()
-    baseDoubleResult.zip(rotatedDoubleResult).foreach {
-      case (br: Double, rr: Double) =>
-        assert(Math.abs(br - rr) < 1e-9, "Predicted double label not the same in rotated learner.")
-    }
-
-    // Check categorical labels are close
-    val baseCatResult = baseCatModel.transform(inputs).getExpected()
-    val rotatedCatResult = rotatedCatModel.transform(inputs).getExpected()
-    val baseF1 = ClassificationMetrics.f1scores(baseCatResult, catLabel)
-    val rotatedF1 = ClassificationMetrics.f1scores(rotatedCatResult, catLabel)
-    // rotatedF1 appears to come in substantially lower than baseF1; this is just a rough sanity check / smoketest.
-    assert(Math.abs(baseF1 - rotatedF1) < 0.15, s"baseF1 ${baseF1} and rotatedF1 ${rotatedF1} are too dissimilar.")
-
-    // Check that rotation features and rotation matrices are the same for both tasks
-    val U = rotatedDoubleModel.trans
-    val rotatedFeatures = rotatedDoubleModel.rotatedFeatures
-    assert(U == rotatedCatModel.trans)
-    assert(rotatedFeatures == rotatedCatModel.rotatedFeatures)
-
-    val rotatedInputs = FeatureRotator.applyRotation(data.map(_._1), rotatedFeatures, U)
-    assert(
-      rotatedDoubleModel.baseModel.transform(rotatedInputs).getExpected().zip(rotatedDoubleResult).count {
-        case (a: Double, b: Double) => Math.abs(a - b) > 1e-9
-      } == 0,
-      "Rotated data passed to base model should map to the same predictions."
-    )
-    assert(
-      rotatedDoubleModel.transform(rotatedInputs).getExpected().zip(baseDoubleResult).count {
-        case (a: Double, b: Double) => Math.abs(a - b) > 1e-9
-      } > 0,
-      "Rotated data passed to rotated model should map to different predictions."
-    )
-
-    val doublyRotatedF1 =
-      ClassificationMetrics.f1scores(rotatedCatModel.transform(rotatedInputs).getExpected(), catLabel)
-    val rotatedF1Base = ClassificationMetrics.f1scores(baseCatModel.transform(rotatedInputs).getExpected(), catLabel)
-
-    // Check that categorical labels change when we feed in different data.
-    assert(
-      rotatedCatModel
-        .transform(rotatedInputs)
-        .getExpected()
-        .zip(baseCatModel.transform(rotatedInputs).getExpected())
-        .count {
-          case (a: Boolean, b: Boolean) => a != b
-        } > 0,
-      s"Categorical labels should substantially change when we feed in different inputs."
-    )
-    // Check that categorical labels are the same as feeding rotated data into base learner.
-    assert(
-      rotatedCatModel
-        .transform(inputs)
-        .getExpected()
-        .zip(rotatedCatModel.baseModel.transform(rotatedInputs).getExpected)
-        .count {
-          case (a: Boolean, b: Boolean) => a != b
-        } == 0,
-      s"Categorical labels should be identical when we feed rotated inputs into base learner."
-    )
-  }
-}
diff --git a/src/test/scala/io/citrine/lolo/transformers/StandardizerTest.scala b/src/test/scala/io/citrine/lolo/transformers/StandardizerTest.scala
deleted file mode 100644
index 0b112911..00000000
--- a/src/test/scala/io/citrine/lolo/transformers/StandardizerTest.scala
+++ /dev/null
@@ -1,212 +0,0 @@
-package io.citrine.lolo.transformers
-
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
-import io.citrine.lolo.linear.{GuessTheMeanLearner, LinearRegressionLearner}
-import io.citrine.lolo.stats.functions.Friedman
-import io.citrine.lolo.stats.metrics.ClassificationMetrics
-import io.citrine.lolo.trees.classification.ClassificationTreeLearner
-import io.citrine.lolo.trees.multitask.MultiTaskTreeLearner
-import org.junit.Test
-
-/**
-  * Created by maxhutch on 2/19/17.
-  */
-@Test
-class StandardizerTest extends SeedRandomMixIn {
-
-  val data: Vector[(Vector[Double], Double)] =
-    TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
-  val weights: Vector[Double] = Vector.fill(data.size)(if (rng.nextBoolean()) rng.nextDouble() else 0.0)
-
-  // Creating another dataset which has 1 feature that has 0 variance.
-  val dataWithConstant: Vector[(Vector[Double], Double)] = data.map(d => (0.0 +: d._1, d._2))
-
-  @Test
-  def testStandardMeanAndVariance(): Unit = {
-    val inputs = data.map(_._1)
-    val inputTransforms = Standardizer.getMultiStandardization(inputs)
-
-    val standardInputs = Standardizer.applyStandardization(inputs, inputTransforms)
-    standardInputs.head.indices.foreach { i =>
-      val values = standardInputs.map(_(i).asInstanceOf[Double])
-      val mean = values.sum / values.size
-      assert(mean < 1.0e-9, s"Standard input ${i} has non-zero mean ${mean}")
-
-      val variance = values.map(Math.pow(_, 2)).sum / values.size
-      assert(Math.abs(variance - 1.0) < 1.0e-9, s"Standard input ${i} has non-unit variance ${variance}")
-    }
-
-    val outputs = data.map(_._2.asInstanceOf[Double])
-    val outputTransform = Standardizer.getStandardization(outputs)
-    val standardOutputs = Standardizer.applyStandardization(outputs, Some(outputTransform)).asInstanceOf[Seq[Double]]
-
-    val mean = standardOutputs.sum / standardOutputs.size
-    assert(mean < 1.0e-9, s"Standard output has non-zero mean ${mean}")
-
-    val variance = standardOutputs.map(Math.pow(_, 2)).sum / standardOutputs.size
-    assert(Math.abs(variance - 1.0) < 1.0e-9, s"Standard output has non-unit variance ${variance}")
-  }
-
-  /**
-    * Guess the mean should be invariant under standardization
-    */
-  @Test
-  def testStandardGTM(): Unit = {
-    val learner = GuessTheMeanLearner()
-    val model = learner.train(data, rng = rng).getModel()
-    val result = model.transform(data.map(_._1)).getExpected()
-
-    val standardLearner = Standardizer(GuessTheMeanLearner())
-    val standardModel = standardLearner.train(data, rng = rng).getModel()
-    val standardResult = standardModel.transform(data.map(_._1)).getExpected()
-
-    result.zip(standardResult).foreach {
-      case (free: Double, standard: Double) =>
-        assert(Math.abs(free - standard) < 1.0e-9, s"${free} and ${standard} should be the same")
-    }
-  }
-
-  /**
-    * Linear regression should be invariant under standardization
-    */
-  @Test
-  def testStandardLinear(): Unit = {
-    val learner = LinearRegressionLearner()
-    val model = learner.train(data, Some(weights)).getModel()
-    val result = model.transform(data.map(_._1))
-    val expected = result.getExpected()
-    val gradient = result.getGradient()
-
-    val standardLearner = Standardizer(learner)
-    val standardModel = standardLearner.train(data, Some(weights)).getModel()
-    val standardResult = standardModel.transform(data.map(_._1))
-    val standardExpected = standardResult.getExpected()
-    val standardGradient = standardResult.getGradient()
-
-    expected.zip(standardExpected).foreach {
-      case (free: Double, standard: Double) =>
-        assert(Math.abs(free - standard) < 1.0e-9, s"${free} and ${standard} should be the same")
-    }
-
-    gradient.get.zip(standardGradient.get).foreach {
-      case (free, standard) =>
-        val diff = free.zip(standard).map { case (f, s) => Math.abs(f - s) }.max
-        assert(diff < 1.0e-9, "Gradients should be the same")
-    }
-  }
-
-  /**
-    * When the variance of a particular feature is 0
-    *
-    * @author Astha Garg
-    */
-  @Test
-  def testStandardWithConstantFeature(): Unit = {
-    val learner = LinearRegressionLearner()
-    val model = learner.train(dataWithConstant, Some(weights)).getModel()
-    val result = model.transform(dataWithConstant.map(_._1))
-    val expected = result.getExpected()
-    val gradient = result.getGradient()
-
-    val standardLearner = Standardizer(learner)
-    val standardModel = standardLearner.train(dataWithConstant, Some(weights)).getModel()
-    val standardResult = standardModel.transform(dataWithConstant.map(_._1))
-    val standardExpected = standardResult.getExpected()
-    val standardGradient = standardResult.getGradient()
-
-    gradient.get.toList.flatten.zip(standardGradient.get.toList.flatten).foreach {
-      case (free: Double, standard: Double) =>
-        assert(
-          Math.abs(free - standard) < 1.0e-9,
-          s"Failed test for gradient. ${free} and ${standard} gradients should be the same"
-        )
-    }
-
-    // The gradient wrt the first constant feature is ill-defined without regularization
-    gradient.get.zip(standardGradient.get).foreach {
-      case (free, standard) =>
-        val diff = free.zip(standard).map { case (f, s) => Math.abs(f - s) }.max
-        assert(diff < 1.0e-9, s"Gradients should be the same. The diff is $diff")
-    }
-
-    // This test fails ~30% of the time when nRows=30, but its not clear why.
-    expected.zip(standardExpected).foreach {
-      case (free: Double, standard: Double) =>
-        assert(
-          Math.abs(free - standard) < 1.0e-9,
-          s"Failed test for expected. ${free} and ${standard} should be the same"
-        )
-    }
-  }
-
-  /**
-    * Ridge regression should depend on standardization
-    */
-  @Test
-  def testStandardRidge(): Unit = {
-    val learner = LinearRegressionLearner(regParam = Some(1.0))
-    val model = learner.train(data).getModel()
-    val result = model.transform(data.map(_._1)).getExpected()
-
-    val standardLearner = new Standardizer(learner)
-    val standardModel = standardLearner.train(data).getModel()
-    val standardResult = standardModel.transform(data.map(_._1)).getExpected()
-
-    result.zip(standardResult).foreach {
-      case (free: Double, standard: Double) =>
-        assert(Math.abs(free - standard) > 1.0e-9, s"${free} and ${standard} should NOT be the same")
-    }
-  }
-
-  @Test
-  def testStandardClassification(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(2048, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      responseBins = Some(2)
-    )
-
-    val learner = ClassificationTreeLearner()
-    val model = learner.train(trainingData).getModel()
-    val result = model.transform(trainingData.map(_._1)).getExpected()
-
-    val standardLearner = Standardizer(learner)
-    val standardModel = standardLearner.train(trainingData).getModel()
-    val standardResult = standardModel.transform(trainingData.map(_._1)).getExpected()
-    result.zip(standardResult).foreach {
-      case (free: String, standard: String) =>
-        assert(free == standard, s"Standard classification tree should be the same")
-    }
-  }
-
-  /**
-    * Test that multitask sandardization has the proper performance and invariants
-    */
-  @Test
-  def testMultiTaskStandardizer(): Unit = {
-    // Generate multi-task training data
-    val (inputs, doubleLabel) = data.unzip
-    val catLabel = inputs.map(x => Friedman.friedmanGrosseSilverman(x) > 15.0)
-
-    // 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.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
-    val baseF1 = ClassificationMetrics.f1scores(baseRes, catLabel)
-    val standardF1 = ClassificationMetrics.f1scores(standardRes, catLabel)
-
-    // Assert some things
-    assert(Math.abs(baseF1 - standardF1) < 1.0e-6, s"Expected training to be invariant the scale of the labels")
-  }
-}
diff --git a/src/test/scala/io/citrine/lolo/transformers/rotator/FeatureRotatorTest.scala b/src/test/scala/io/citrine/lolo/transformers/rotator/FeatureRotatorTest.scala
new file mode 100644
index 00000000..f5b80439
--- /dev/null
+++ b/src/test/scala/io/citrine/lolo/transformers/rotator/FeatureRotatorTest.scala
@@ -0,0 +1,257 @@
+package io.citrine.lolo.transformers.rotator
+
+import breeze.linalg.{det, DenseMatrix}
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.linear.{GuessTheMeanLearner, LinearRegressionLearner}
+import io.citrine.lolo.stats.functions.Friedman
+import io.citrine.lolo.trees.classification.ClassificationTreeLearner
+import io.citrine.lolo.trees.regression.RegressionTreeLearner
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
+import org.junit.Test
+
+/**
+  * Created by gregor-robinson on 2020-01-06.
+  */
+@Test
+class FeatureRotatorTest extends SeedRandomMixIn {
+
+  val data: Seq[TrainingRow[Double]] = DataGenerator
+    .generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+    .withBinnedInputs(bins = Seq((0, 8)))
+    .data
+  val weightedData: Seq[TrainingRow[Double]] = data.map { row =>
+    row.withWeight(if (rng.nextBoolean()) rng.nextDouble() else 0.0)
+  }
+
+  @Test
+  def testRandomRotation(): Unit = {
+    val inputs = data.map(_.inputs)
+    val numRealInputs = inputs.head.length - 1
+    for (i <- 1 to 10) {
+      val U = FeatureRotator.getRandomRotation(numRealInputs)
+      assert(U.rows == numRealInputs)
+      assert(U.cols == numRealInputs)
+
+      // Check that the matrix is unitary.
+      (U.t * U - DenseMatrix.eye[Double](numRealInputs)).toArray.foreach { x =>
+        assert(Math.abs(x) < 1e-9)
+      }
+
+      // Check that the determinant is one.
+      val d = det(U)
+      assert(Math.abs(d - 1.0) < 1e-9, s"Determinant of U should be 1 but is ${d}")
+    }
+  }
+
+  @Test
+  def testApplyRotation(): Unit = {
+    val inputs = data.map(_.inputs)
+    val numRealInputs = inputs.head.length - 1
+    val featuresToRotate = (1 to numRealInputs).asInstanceOf[IndexedSeq[Int]]
+
+    // Create rotation matrix that exchanges first, second, and last real inputs
+    val U = DenseMatrix.eye[Double](dim = numRealInputs)
+    U(0, 0) = 0.0
+    U(1, 1) = 0.0
+    U(numRealInputs - 1, numRealInputs - 1) = 0.0
+    U(0, 1) = 1.0
+    U(1, numRealInputs - 1) = 1.0
+    U(numRealInputs - 1, 0) = 1.0
+
+    val rotatedInputs = FeatureRotator.applyRotation(inputs, featuresToRotate, U)
+    inputs.indices.foreach { i =>
+      assert(inputs(i)(0) == rotatedInputs(i)(0), "Failed to leave categorical invariant.")
+      assert(
+        Math.abs(inputs(i)(1).asInstanceOf[Double] - rotatedInputs(i)(numRealInputs).asInstanceOf[Double]) < 1e-9,
+        "Failed to exchange coordinates."
+      )
+      assert(
+        Math.abs(inputs(i)(2).asInstanceOf[Double] - rotatedInputs(i)(1).asInstanceOf[Double]) < 1e-9,
+        "Failed to exchange coordinates."
+      )
+      assert(
+        Math.abs(inputs(i)(numRealInputs).asInstanceOf[Double] - rotatedInputs(i)(2).asInstanceOf[Double]) < 1e-9,
+        "Failed to exchange coordinates."
+      )
+      (3 to numRealInputs - 1).foreach { j =>
+        assert(
+          Math.abs(inputs(i)(j).asInstanceOf[Double] - rotatedInputs(i)(j).asInstanceOf[Double]) < 1e-9,
+          "Unexpected modification of un-rotated coordinate."
+        )
+      }
+    }
+
+    // Check that we can undo rotation by applying the transpose
+    val unrotatedInputs = FeatureRotator.applyRotation(rotatedInputs, featuresToRotate, U.t)
+    inputs.indices.foreach { i =>
+      assert(inputs(i)(0) == unrotatedInputs(i)(0), "Failed to leave categorical invariant.")
+      (1 to numRealInputs - 1).foreach { j =>
+        assert(
+          Math.abs(inputs(i)(j).asInstanceOf[Double] - unrotatedInputs(i)(j).asInstanceOf[Double]) < 1e-9,
+          "Unexpected modification of coordinate after applying transpose."
+        )
+      }
+    }
+  }
+
+  /**
+    * Functions getLoss and getPredictedVsActual should pass through to base learner.
+    */
+  @Test
+  def testPassthroughFunctions(): Unit = {
+    val rotatedTrainingResult = FeatureRotator(RegressionTreeLearner()).train(data, rng = rng)
+
+    assert(
+      rotatedTrainingResult.loss == rotatedTrainingResult.baseTrainingResult.loss,
+      "Function getLoss() should pass through to base learner."
+    )
+
+    rotatedTrainingResult.predictedVsActual.foreach { x =>
+      x.zip(data).foreach {
+        case (a, b) =>
+          assert(a._1 == b.inputs, "getPredictedVsActual must return the correct training inputs.")
+          assert(a._2 == b.label, "getPredictedVsActual must return the correct predicted value.")
+          assert(a._3 == b.label, "getPredictedVsActual must return the correct actual value.")
+      }
+    }
+  }
+
+  /**
+    * Guess the mean should be invariant under rotation.
+    */
+  @Test
+  def testRotatedGTM(): Unit = {
+    val learner = GuessTheMeanLearner()
+    val model = learner.train(data, rng = rng).model
+    val result = model.transform(data.map(_.inputs)).expected
+
+    val rotatedLearner = FeatureRotator(GuessTheMeanLearner())
+    val rotatedModel = rotatedLearner.train(data, rng = rng).model
+    val rotatedResult = rotatedModel.transform(data.map(_.inputs)).expected
+
+    result.zip(rotatedResult).foreach {
+      case (free: Double, rotated: Double) =>
+        assert(Math.abs(free - rotated) < 1.0e-9, s"${free} and ${rotated} should be the same")
+    }
+  }
+
+  /**
+    * Linear regression should be invariant under rotation
+    */
+  @Test
+  def testRotatedLinear(): Unit = {
+    val learner = LinearRegressionLearner()
+    val model = learner.train(weightedData, rng).model
+    val result = model.transform(data.map(_.inputs))
+    val expected = result.expected
+    val gradient = result.gradient
+
+    val rotatedLearner = FeatureRotator(learner)
+    val rotatedModel = rotatedLearner.train(weightedData, rng).model
+    val rotatedResult = rotatedModel.transform(data.map(_.inputs))
+    val rotatedExpected = rotatedResult.expected
+    val rotatedGradient = rotatedResult.gradient
+
+    expected.zip(rotatedExpected).foreach {
+      case (free: Double, rotated: Double) =>
+        assert(Math.abs(free - rotated) < 1.0e-9, s"${free} and ${rotated} should be the same")
+    }
+
+    gradient.get.zip(rotatedGradient.get).foreach {
+      case (free, rotated) =>
+        val diff = free.zip(rotated).map { case (f, s) => Math.abs(f - s) }.max
+        assert(diff < 1.0e-9, "Gradients should be the same")
+    }
+  }
+
+  /**
+    * Ridge regression should not depend on rotation
+    */
+  @Test
+  def testRotatedRidge(): Unit = {
+    val learner = LinearRegressionLearner(regParam = Some(1.0))
+    val model = learner.train(data, rng = rng).model
+    val result = model.transform(data.map(_.inputs)).expected
+
+    val rotatedLearner = new FeatureRotator(learner)
+    val rotatedModel = rotatedLearner.train(data, rng = rng).model
+    val rotatedResult = rotatedModel.transform(data.map(_.inputs)).expected
+
+    result.zip(rotatedResult).foreach {
+      case (free: Double, rotated: Double) =>
+        assert(Math.abs(free - rotated) < 1.0e-9, s"${free} and ${rotated} should be the same")
+    }
+  }
+
+  /**
+    * Verify expected properties of rotated-input regression trees
+    */
+  @Test
+  def testRotatedRegressionTree(): Unit = {
+    val learner = RegressionTreeLearner()
+    val model = learner.train(data, rng = rng).model
+    val result = model.transform(data.map(_.inputs)).expected
+
+    val rotatedLearner = new FeatureRotator(learner)
+    val rotatedModel = rotatedLearner.train(data, rng = rng).model.asInstanceOf[RotatedFeatureModel[Double]]
+    var rotatedResult = rotatedModel.transform(data.map(_.inputs)).expected
+    result.zip(rotatedResult).foreach {
+      case (free: Double, rotated: Double) =>
+        assert(Math.abs(free - rotated) < 1.0e-9, s"${free} and ${rotated} should be the same")
+    }
+
+    val rotatedData = FeatureRotator.applyRotation(data.map(_.inputs), rotatedModel.rotatedFeatures, rotatedModel.trans)
+    rotatedResult = rotatedModel.transform(rotatedData).expected
+    // Check that labels change when we feed in different data.
+    assert(
+      rotatedResult.zip(result).map { case (a: Double, b: Double) => a - b }.count { x => Math.abs(x) > 1e-9 } > 0,
+      "Rotated data passed to rotated model should not map to the same predictions."
+    )
+    val baseResult = rotatedModel.baseModel.transform(rotatedData).expected
+    // Check that labels are the same as feeding rotated data into base learner.
+    assert(
+      baseResult.zip(result).map { case (a: Double, b: Double) => a - b }.count { x => Math.abs(x) > 1e-9 } == 0,
+      "Rotated data passed to base model should map to the same predictions."
+    )
+  }
+
+  /**
+    * Verify expected properties of rotated-input classification trees
+    */
+  @Test
+  def testRotatedClassificationTree(): Unit = {
+    val classificationData = DataGenerator
+      .generate(2048, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedLabels(bins = 2)
+      .data
+
+    val learner = ClassificationTreeLearner()
+    val model = learner.train(classificationData, rng = rng).model
+    val result = model.transform(classificationData.map(_.inputs)).expected
+
+    val rotatedLearner = new FeatureRotator(learner)
+    val rotatedModel =
+      rotatedLearner.train(classificationData, rng = rng).model.asInstanceOf[RotatedFeatureModel[String]]
+    var rotatedResult = rotatedModel.transform(classificationData.map(_.inputs)).expected
+
+    result.zip(rotatedResult).foreach {
+      case (free: Any, rotated: Any) =>
+        assert(free == rotated, s"${free} and ${rotated} should be the same")
+    }
+
+    val rotatedData =
+      FeatureRotator.applyRotation(classificationData.map(_.inputs), rotatedModel.rotatedFeatures, rotatedModel.trans)
+    rotatedResult = rotatedModel.transform(rotatedData).expected
+    // Check that labels change when we feed in different data.
+    assert(
+      rotatedResult.zip(result).count { case (a: String, b: String) => a != b } > 0,
+      "Rotated data passed to rotated model should not map to the same predictions."
+    )
+    val baseResult = rotatedModel.baseModel.transform(rotatedData).expected
+    // Check that labels are the same as feeding rotated data into base learner.
+    assert(
+      baseResult.zip(result).count { case (a: String, b: String) => a != b } == 0,
+      s"Categorical labels should be identical when we feed rotated inputs into base learner."
+    )
+  }
+}
diff --git a/src/test/scala/io/citrine/lolo/transformers/rotator/MultiTaskFeatureRotatorTest.scala b/src/test/scala/io/citrine/lolo/transformers/rotator/MultiTaskFeatureRotatorTest.scala
new file mode 100644
index 00000000..5403a9ea
--- /dev/null
+++ b/src/test/scala/io/citrine/lolo/transformers/rotator/MultiTaskFeatureRotatorTest.scala
@@ -0,0 +1,100 @@
+package io.citrine.lolo.transformers.rotator
+
+import io.citrine.lolo.stats.functions.Friedman
+import io.citrine.lolo.stats.metrics.ClassificationMetrics
+import io.citrine.lolo.trees.multitask.MultiTaskTreeLearner
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
+import org.junit.Test
+
+class MultiTaskFeatureRotatorTest extends SeedRandomMixIn {
+
+  /**
+    * Verify that rotated-input multi-task trees have expected properties
+    */
+  @Test
+  def testMultiTaskRotator(): Unit = {
+    val rawData = DataGenerator
+      .generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .data
+
+    // Generate multi-task training data
+    val inputs = rawData.map(_.inputs)
+    val doubleLabel = rawData.map(_.label)
+    val catLabel = inputs.map(x => Friedman.friedmanGrosseSilverman(x.asInstanceOf[Seq[Double]]) > 15.0)
+
+    // Sparsify the categorical labels
+    val sparseCatLabel = catLabel.map(x => if (rng.nextBoolean()) x else null)
+    val labels = Vector(doubleLabel, sparseCatLabel).transpose
+    val multiTaskData = rawData.zip(labels).map { case (row, label) => row.withLabel(label) }
+
+    // Train and evaluate rotated models on original and rotated features
+    val baseLearner = MultiTaskTreeLearner()
+    val rotatedLearner = MultiTaskFeatureRotator(MultiTaskTreeLearner())
+
+    val baseTrainingResult = baseLearner.train(multiTaskData, rng = rng)
+    val baseDoubleModel = baseTrainingResult.models.head
+    val baseCatModel = baseTrainingResult.models.last
+    val rotatedTrainingResult = rotatedLearner.train(multiTaskData, rng = rng)
+    val rotatedDoubleModel = rotatedTrainingResult.models.head.asInstanceOf[RotatedFeatureModel[Double]]
+    val rotatedCatModel = rotatedTrainingResult.models.last.asInstanceOf[RotatedFeatureModel[Boolean]]
+
+    // Check double labels are the same
+    val baseDoubleResult = baseDoubleModel.transform(inputs).expected
+    val rotatedDoubleResult = rotatedDoubleModel.transform(inputs).expected
+    baseDoubleResult.zip(rotatedDoubleResult).foreach {
+      case (br: Double, rr: Double) =>
+        assert(Math.abs(br - rr) < 1e-9, "Predicted double label not the same in rotated learner.")
+    }
+
+    // Check categorical labels are close
+    val baseCatResult = baseCatModel.transform(inputs).expected
+    val rotatedCatResult = rotatedCatModel.transform(inputs).expected
+    val baseF1 = ClassificationMetrics.f1scores(baseCatResult, catLabel)
+    val rotatedF1 = ClassificationMetrics.f1scores(rotatedCatResult, catLabel)
+    // rotatedF1 appears to come in substantially lower than baseF1; this is just a rough sanity check / smoketest.
+    assert(Math.abs(baseF1 - rotatedF1) < 0.15, s"baseF1 ${baseF1} and rotatedF1 ${rotatedF1} are too dissimilar.")
+
+    // Check that rotation features and rotation matrices are the same for both tasks
+    val U = rotatedDoubleModel.trans
+    val rotatedFeatures = rotatedDoubleModel.rotatedFeatures
+    assert(U == rotatedCatModel.trans)
+    assert(rotatedFeatures == rotatedCatModel.rotatedFeatures)
+
+    val rotatedInputs = FeatureRotator.applyRotation(inputs, rotatedFeatures, U)
+    assert(
+      rotatedDoubleModel.baseModel.transform(rotatedInputs).expected.zip(rotatedDoubleResult).count {
+        case (a: Double, b: Double) => Math.abs(a - b) > 1e-9
+      } == 0,
+      "Rotated data passed to base model should map to the same predictions."
+    )
+    assert(
+      rotatedDoubleModel.transform(rotatedInputs).expected.zip(baseDoubleResult).count {
+        case (a: Double, b: Double) => Math.abs(a - b) > 1e-9
+      } > 0,
+      "Rotated data passed to rotated model should map to different predictions."
+    )
+
+    // Check that categorical labels change when we feed in different data.
+    assert(
+      rotatedCatModel
+        .transform(rotatedInputs)
+        .expected
+        .zip(baseCatModel.transform(rotatedInputs).expected)
+        .count {
+          case (a: Boolean, b: Boolean) => a != b
+        } > 0,
+      s"Categorical labels should substantially change when we feed in different inputs."
+    )
+    // Check that categorical labels are the same as feeding rotated data into base learner.
+    assert(
+      rotatedCatModel
+        .transform(inputs)
+        .expected
+        .zip(rotatedCatModel.baseModel.transform(rotatedInputs).expected)
+        .count {
+          case (a: Boolean, b: Boolean) => a != b
+        } == 0,
+      s"Categorical labels should be identical when we feed rotated inputs into base learner."
+    )
+  }
+}
diff --git a/src/test/scala/io/citrine/lolo/transformers/standardizer/MultiTaskStandardizerTest.scala b/src/test/scala/io/citrine/lolo/transformers/standardizer/MultiTaskStandardizerTest.scala
new file mode 100644
index 00000000..c64333da
--- /dev/null
+++ b/src/test/scala/io/citrine/lolo/transformers/standardizer/MultiTaskStandardizerTest.scala
@@ -0,0 +1,52 @@
+package io.citrine.lolo.transformers.standardizer
+
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.stats.functions.Friedman
+import io.citrine.lolo.stats.metrics.ClassificationMetrics
+import io.citrine.lolo.trees.multitask.MultiTaskTreeLearner
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
+import org.junit.Test
+
+class MultiTaskStandardizerTest extends SeedRandomMixIn {
+
+  val data: Seq[TrainingRow[Double]] =
+    DataGenerator.generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng).data
+
+  /**
+    * Test that multitask sandardization has the proper performance and invariants
+    */
+  @Test
+  def testMultiTaskStandardizer(): Unit = {
+    // Generate multi-task training data
+    val (inputs, doubleLabel) = data.map(row => (row.inputs, row.label)).unzip
+    val catLabel = inputs.map(x => Friedman.friedmanGrosseSilverman(x.asInstanceOf[Seq[Double]]) > 15.0)
+
+    // Sparsify the categorical labels
+    val sparseCatLabel = catLabel.map(x => if (rng.nextBoolean()) x else null)
+    val labels = Vector(doubleLabel, sparseCatLabel).transpose
+    val baseRows = data.zip(labels).map {
+      case (row, label) => row.withLabel(label)
+    }
+
+    // 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) }
+    val rescaledRows = data.zip(rescaledLabels).map {
+      case (row, label) => row.withLabel(label)
+    }
+
+    // Train and evaluate standard models on original and rescaled labels
+    val standardizer = MultiTaskStandardizer(MultiTaskTreeLearner())
+    val baseRes = standardizer.train(baseRows).models.last.transform(inputs).expected
+    val standardRes = standardizer.train(rescaledRows).models.last.transform(inputs).expected
+    // Train and evaluate unstandardized model on rescaled labels
+
+    // Compute metrics for each of the models
+    val baseF1 = ClassificationMetrics.f1scores(baseRes, catLabel)
+    val standardF1 = ClassificationMetrics.f1scores(standardRes, catLabel)
+
+    // Assert some things
+    assert(Math.abs(baseF1 - standardF1) < 1.0e-6, s"Expected training to be invariant the scale of the labels")
+  }
+}
diff --git a/src/test/scala/io/citrine/lolo/transformers/standardizer/StandardizerTest.scala b/src/test/scala/io/citrine/lolo/transformers/standardizer/StandardizerTest.scala
new file mode 100644
index 00000000..40b5e5a1
--- /dev/null
+++ b/src/test/scala/io/citrine/lolo/transformers/standardizer/StandardizerTest.scala
@@ -0,0 +1,179 @@
+package io.citrine.lolo.transformers.standardizer
+
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.linear.{GuessTheMeanLearner, LinearRegressionLearner}
+import io.citrine.lolo.stats.functions.Friedman
+import io.citrine.lolo.trees.classification.ClassificationTreeLearner
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
+import org.junit.Test
+
+@Test
+class StandardizerTest extends SeedRandomMixIn {
+
+  val data: Seq[TrainingRow[Double]] =
+    DataGenerator.generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng).data
+  val weights: Vector[Double] = Vector.fill(data.size)(if (rng.nextBoolean()) rng.nextDouble() else 0.0)
+  // Creating another dataset which has 1 feature that has 0 variance.
+  val dataWithConstant: Seq[TrainingRow[Double]] = data.map(_.mapInputs(x => 0 +: x))
+
+  val weightedData: Seq[TrainingRow[Double]] = data.zip(weights).map { case (row, w) => row.withWeight(w) }
+  val weightedDataWithConstant: Seq[TrainingRow[Double]] =
+    dataWithConstant.zip(weights).map { case (row, w) => row.withWeight(w) }
+
+  @Test
+  def testStandardMeanAndVariance(): Unit = {
+    val inputs = data.map(_.inputs)
+    val inputTransforms = Standardization.buildMulti(inputs)
+
+    val standardInputs = inputs.map(input => Standardization.applyMulti(input, inputTransforms))
+    standardInputs.head.indices.foreach { i =>
+      val values = standardInputs.map(_(i).asInstanceOf[Double])
+      val mean = values.sum / values.size
+      assert(mean < 1.0e-9, s"Standard input ${i} has non-zero mean ${mean}")
+
+      val variance = values.map(Math.pow(_, 2)).sum / values.size
+      assert(Math.abs(variance - 1.0) < 1.0e-9, s"Standard input ${i} has non-unit variance ${variance}")
+    }
+
+    val outputs = data.map(_.label)
+    val outputTransform = Standardization.build(outputs)
+    val standardOutputs = outputs.map(outputTransform.apply)
+
+    val mean = standardOutputs.sum / standardOutputs.size
+    assert(mean < 1.0e-9, s"Standard output has non-zero mean ${mean}")
+
+    val variance = standardOutputs.map(Math.pow(_, 2)).sum / standardOutputs.size
+    assert(Math.abs(variance - 1.0) < 1.0e-9, s"Standard output has non-unit variance ${variance}")
+  }
+
+  /**
+    * Guess the mean should be invariant under standardization
+    */
+  @Test
+  def testStandardGTM(): Unit = {
+    val learner = GuessTheMeanLearner()
+    val model = learner.train(data, rng = rng).model
+    val result = model.transform(data.map(_.inputs)).expected
+
+    val standardLearner = RegressionStandardizer(GuessTheMeanLearner())
+    val standardModel = standardLearner.train(data, rng = rng).model
+    val standardResult = standardModel.transform(data.map(_.inputs)).expected
+
+    result.zip(standardResult).foreach {
+      case (free: Double, standard: Double) =>
+        assert(Math.abs(free - standard) < 1.0e-9, s"${free} and ${standard} should be the same")
+    }
+  }
+
+  /**
+    * Linear regression should be invariant under standardization
+    */
+  @Test
+  def testStandardLinear(): Unit = {
+    val learner = LinearRegressionLearner()
+    val model = learner.train(weightedData).model
+    val result = model.transform(data.map(_.inputs))
+    val expected = result.expected
+    val gradient = result.gradient
+
+    val standardLearner = RegressionStandardizer(learner)
+    val standardModel = standardLearner.train(weightedData).model
+    val standardResult = standardModel.transform(data.map(_.inputs))
+    val standardExpected = standardResult.expected
+    val standardGradient = standardResult.gradient
+
+    expected.zip(standardExpected).foreach {
+      case (free: Double, standard: Double) =>
+        assert(Math.abs(free - standard) < 1.0e-9, s"${free} and ${standard} should be the same")
+    }
+
+    gradient.get.zip(standardGradient.get).foreach {
+      case (free, standard) =>
+        val diff = free.zip(standard).map { case (f, s) => Math.abs(f - s) }.max
+        assert(diff < 1.0e-9, "Gradients should be the same")
+    }
+  }
+
+  /**
+    * When the variance of a particular feature is 0
+    *
+    * @author Astha Garg
+    */
+  @Test
+  def testStandardWithConstantFeature(): Unit = {
+    val learner = LinearRegressionLearner()
+    val model = learner.train(weightedDataWithConstant).model
+    val result = model.transform(dataWithConstant.map(_.inputs))
+    val expected = result.expected
+    val gradient = result.gradient
+
+    val standardLearner = RegressionStandardizer(learner)
+    val standardModel = standardLearner.train(weightedDataWithConstant).model
+    val standardResult = standardModel.transform(dataWithConstant.map(_.inputs))
+    val standardExpected = standardResult.expected
+    val standardGradient = standardResult.gradient
+
+    gradient.get.toList.flatten.zip(standardGradient.get.toList.flatten).foreach {
+      case (free: Double, standard: Double) =>
+        assert(
+          Math.abs(free - standard) < 1.0e-9,
+          s"Failed test for gradient. ${free} and ${standard} gradients should be the same"
+        )
+    }
+
+    // The gradient wrt the first constant feature is ill-defined without regularization
+    gradient.get.zip(standardGradient.get).foreach {
+      case (free, standard) =>
+        val diff = free.zip(standard).map { case (f, s) => Math.abs(f - s) }.max
+        assert(diff < 1.0e-9, s"Gradients should be the same. The diff is $diff")
+    }
+
+    // This test fails ~30% of the time when nRows=30, but its not clear why.
+    expected.zip(standardExpected).foreach {
+      case (free: Double, standard: Double) =>
+        assert(
+          Math.abs(free - standard) < 1.0e-9,
+          s"Failed test for expected. ${free} and ${standard} should be the same"
+        )
+    }
+  }
+
+  /**
+    * Ridge regression should depend on standardization
+    */
+  @Test
+  def testStandardRidge(): Unit = {
+    val learner = LinearRegressionLearner(regParam = Some(1.0))
+    val model = learner.train(data).model
+    val result = model.transform(data.map(_.inputs)).expected
+
+    val standardLearner = RegressionStandardizer(learner)
+    val standardModel = standardLearner.train(data).model
+    val standardResult = standardModel.transform(data.map(_.inputs)).expected
+
+    result.zip(standardResult).foreach {
+      case (free: Double, standard: Double) =>
+        assert(Math.abs(free - standard) > 1.0e-9, s"${free} and ${standard} should NOT be the same")
+    }
+  }
+
+  @Test
+  def testStandardClassification(): Unit = {
+    val trainingData = DataGenerator
+      .generate(2048, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedLabels(bins = 2)
+      .data
+
+    val learner = ClassificationTreeLearner()
+    val model = learner.train(trainingData).model
+    val result = model.transform(trainingData.map(_.inputs)).expected
+
+    val standardLearner = ClassificationStandardizer(learner)
+    val standardModel = standardLearner.train(trainingData).model
+    val standardResult = standardModel.transform(trainingData.map(_.inputs)).expected
+    result.zip(standardResult).foreach {
+      case (free: String, standard: String) =>
+        assert(free == standard, s"Standard classification tree should be the same")
+    }
+  }
+}
diff --git a/src/test/scala/io/citrine/lolo/trees/classification/ClassificationTreeTest.scala b/src/test/scala/io/citrine/lolo/trees/classification/ClassificationTreeTest.scala
index f29604b4..ffd182ec 100644
--- a/src/test/scala/io/citrine/lolo/trees/classification/ClassificationTreeTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/classification/ClassificationTreeTest.scala
@@ -1,13 +1,11 @@
 package io.citrine.lolo.trees.classification
 
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
 import io.citrine.lolo.stats.functions.Friedman
 import org.junit.Test
 import org.scalatest.Assertions._
 
-/**
-  * Created by maxhutch on 12/2/16.
-  */
 @Test
 class ClassificationTreeTest extends SeedRandomMixIn {
 
@@ -18,36 +16,36 @@ class ClassificationTreeTest extends SeedRandomMixIn {
   def testFeatureImportanceNaN(): Unit = {
     val X = Vector.fill(100) {
       val input = Vector.fill(10)(1.0)
-      (input, 2.0)
+      TrainingRow(input, 2.0)
     }
 
-    val DTLearner = new ClassificationTreeLearner()
+    val DTLearner = ClassificationTreeLearner()
     val DTMeta = DTLearner.train(X)
-    val DT = DTMeta.getModel()
-    assert(DTMeta.getFeatureImportance().get.forall(v => !v.isNaN))
+    assert(DTMeta.featureImportance.get.forall(v => !v.isNaN))
   }
 
   @Test
   def testBinary(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils.generateTrainingData(2048, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      responseBins = Some(2)
-    )
+    val trainingData = DataGenerator
+      .generate(2048, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedLabels(bins = 2)
+      .data
+
     val DTLearner = ClassificationTreeLearner()
     val DTMeta = DTLearner.train(trainingData)
-    val DT = DTMeta.getModel()
+    val DT = DTMeta.model
 
     /* We should be able to memorize the inputs */
-    val output = DT.transform(trainingData.map(_._1))
-    trainingData.zip(output.getExpected()).foreach {
-      case ((x, a), p) =>
-        assert(a == p, s"${a} != ${p} for ${x}")
+    val output = DT.transform(trainingData.map(_.inputs))
+    trainingData.zip(output.expected).foreach {
+      case (row, p) =>
+        assert(row.label == p, s"${row.label} != $p for ${row.inputs}")
     }
-    assert(output.getGradient().isEmpty)
-    output.getDepth().foreach(d => assert(d > 0))
+    assert(output.gradient.isEmpty)
+    output.depth.foreach(d => assert(d > 0))
 
     /* The first features should be the most important */
-    val importances = DTMeta.getFeatureImportance().get
+    val importances = DTMeta.featureImportance.get
     assert(importances.slice(0, 5).min > importances.slice(5, importances.size).max)
   }
 
@@ -56,31 +54,31 @@ class ClassificationTreeTest extends SeedRandomMixIn {
     */
   @Test
   def longerTest(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils
-        .generateTrainingData(1024, 12, noise = 0.05, function = Friedman.friedmanSilverman, rng = rng),
-      responseBins = Some(16)
-    )
+    val trainingData = DataGenerator
+      .generate(1024, 12, noise = 0.05, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedLabels(bins = 16)
+      .data
+
     val DTLearner = ClassificationTreeLearner()
     val N = 100
     val start = System.nanoTime()
     val DTMeta = DTLearner.train(trainingData)
-    val DT = DTMeta.getModel()
+    val DT = DTMeta.model
     (0 until N).map(i => DTLearner.train(trainingData))
     val duration = (System.nanoTime() - start) / 1.0e9
 
     println(s"Training large case took ${duration / N} s")
 
     /* We should be able to memorize the inputs */
-    val output = DT.transform(trainingData.map(_._1))
-    trainingData.zip(output.getExpected()).foreach {
-      case ((x, a), p) => assert(a == p, s"${a} != ${p} for ${x}")
+    val output = DT.transform(trainingData.map(_.inputs))
+    trainingData.zip(output.expected).foreach {
+      case (row, p) => assert(row.label == p, s"${row.label} != $p for ${row.inputs}")
     }
-    assert(output.getGradient().isEmpty)
-    output.getDepth().foreach(d => assert(d > 4 && d < 17, s"Depth is ${d}"))
+    assert(output.gradient.isEmpty)
+    output.depth.foreach(d => assert(d > 4 && d < 17, s"Depth is ${d}"))
 
     /* The first feature should be the most important */
-    val importances = DTMeta.getFeatureImportance().get
+    val importances = DTMeta.featureImportance.get
     assert(importances.slice(0, 5).min > importances.slice(5, importances.size).max)
   }
 
@@ -89,29 +87,29 @@ class ClassificationTreeTest extends SeedRandomMixIn {
     */
   @Test
   def testCategorical(): Unit = {
-    val trainingData = TestUtils.binTrainingData(
-      TestUtils
-        .generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-      inputBins = Seq((0, 8)),
-      responseBins = Some(16)
-    )
+    val trainingData = DataGenerator
+      .generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .withBinnedLabels(bins = 16)
+      .data
+
     val DTLearner = ClassificationTreeLearner()
     val N = 100
     val start = System.nanoTime()
-    val DT = DTLearner.train(trainingData).getModel()
+    val DT = DTLearner.train(trainingData).model
     (0 until N).map(i => DTLearner.train(trainingData))
     val duration = (System.nanoTime() - start) / 1.0e9
 
     println(s"Training large case took ${duration / N} s")
 
     /* We should be able to memorize the inputs */
-    val output = DT.transform(trainingData.map(_._1))
-    trainingData.zip(output.getExpected()).foreach {
-      case ((x, a), p) =>
-        assert(a == p)
+    val output = DT.transform(trainingData.map(_.inputs))
+    trainingData.zip(output.expected).foreach {
+      case (row, p) =>
+        assert(row.label == p)
     }
-    assert(output.getGradient().isEmpty)
-    output.getDepth().foreach(d => assert(d > 3 && d < 18, s"Depth is ${d}"))
+    assert(output.gradient.isEmpty)
+    output.depth.foreach(d => assert(d > 3 && d < 18, s"Depth is ${d}"))
   }
 }
 
diff --git a/src/test/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculatorTest.scala b/src/test/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculatorTest.scala
index a8f78299..8e53aac5 100644
--- a/src/test/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculatorTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculatorTest.scala
@@ -2,9 +2,6 @@ package io.citrine.lolo.trees.impurity
 
 import org.junit.Test
 
-/**
-  * Created by maxhutch on 12/1/16.
-  */
 @Test
 class MultiImpurityCalculatorTest {
 
@@ -14,9 +11,9 @@ class MultiImpurityCalculatorTest {
   @Test
   def testImpurity(): Unit = {
     val data = Seq(
-      (Array[AnyVal](0.0, 1.toChar), 1.0),
-      (Array[AnyVal](1.0, 2.toChar), 1.0),
-      (Array[AnyVal](100.0, 3.toChar), 0.0)
+      (Vector[AnyVal](0.0, 1.toChar), 1.0),
+      (Vector[AnyVal](1.0, 2.toChar), 1.0),
+      (Vector[AnyVal](100.0, 3.toChar), 0.0)
     )
     val calculator = MultiImpurityCalculator.build(data.map(_._1), data.map(_._2))
 
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 f71b1e34..f9cf3875 100644
--- a/src/test/scala/io/citrine/lolo/trees/multitask/MultiTaskTreeTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/multitask/MultiTaskTreeTest.scala
@@ -1,6 +1,7 @@
 package io.citrine.lolo.trees.multitask
 
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
 import io.citrine.lolo.stats.functions.Friedman
 import io.citrine.lolo.stats.metrics.ClassificationMetrics
 import io.citrine.lolo.trees.classification.{ClassificationTree, ClassificationTreeLearner}
@@ -13,11 +14,12 @@ import org.scalatest.Assertions._
 class MultiTaskTreeTest extends SeedRandomMixIn {
 
   // Set up some data.
-  val raw: Seq[(Vector[Double], Double)] =
-    TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
-  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
+  val realRows: Seq[TrainingRow[Double]] =
+    DataGenerator.generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng).data
+  val (inputs: Seq[Vector[Double]], realLabel: Seq[Double]) = realRows.map(row => (row.inputs, row.label)).unzip
+  val catLabel: Seq[Boolean] = realRows.map(_.label > realLabel.max / 2.0)
+  val labels: Seq[Vector[Any]] = Vector(realLabel, catLabel).transpose
+  val multiRows: Seq[TrainingRow[Vector[Any]]] = realRows.zip(labels).map { case (row, label) => row.withLabel(label) }
 
   /**
     * Test that the tree memorizes the data when it is complete
@@ -25,12 +27,12 @@ class MultiTaskTreeTest extends SeedRandomMixIn {
   @Test
   def testTwoLabels(): Unit = {
     val learner = MultiTaskTreeLearner()
-    val models = learner.train(inputs.zip(labels), rng = rng).getModels()
+    val models = learner.train(multiRows, rng = rng).models
     assert(models.size == 2)
     assert(models.head.isInstanceOf[RegressionTree])
     assert(models.last.isInstanceOf[ClassificationTree])
-    val realResults = models.head.transform(inputs).getExpected().asInstanceOf[Seq[Double]]
-    val catResults = models.last.transform(inputs).getExpected().asInstanceOf[Seq[Boolean]]
+    val realResults = models.head.transform(inputs).expected.asInstanceOf[Seq[Double]]
+    val catResults = models.last.transform(inputs).expected.asInstanceOf[Seq[Boolean]]
 
     assert(realLabel.zip(realResults).forall(p => p._1 == p._2))
     assert(catLabel.zip(catResults).forall(p => p._1 == p._2))
@@ -49,11 +51,12 @@ class MultiTaskTreeTest extends SeedRandomMixIn {
       }
     )
     val sparseLabels = Vector(realLabel, sparseCat).transpose
+    val sparseRows = TrainingRow.build(inputs.zip(sparseLabels))
 
     val learner = MultiTaskTreeLearner()
-    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]]
+    val models = learner.train(sparseRows).models
+    val realResults = models.head.transform(inputs).expected.asInstanceOf[Seq[Double]]
+    val catResults = models.last.transform(inputs).expected.asInstanceOf[Seq[Boolean]]
 
     assert(realLabel.zip(realResults).forall(p => p._1 == p._2))
     assert(sparseCat.zip(catResults).forall(p => p._1 == null || p._1 == p._2))
@@ -73,15 +76,18 @@ class MultiTaskTreeTest extends SeedRandomMixIn {
     )
     val sparseLabels = Vector(realLabel, sparseCat).transpose
 
+    val sparseCatRows = TrainingRow.build(inputs.zip(sparseCat))
+    val sparseMultiRows = TrainingRow.build(inputs.zip(sparseLabels)).filterNot(_.label == null)
+
     val learner = MultiTaskTreeLearner()
-    val models = learner.train(inputs.zip(sparseLabels), rng = rng).getModels()
-    val catResults = models.last.transform(inputs).getExpected().asInstanceOf[Seq[Boolean]]
+    val models = learner.train(sparseMultiRows, rng = rng).models
+    val catResults = models.last.transform(inputs).expected.asInstanceOf[Seq[Boolean]]
 
     val reference = ClassificationTreeLearner()
-      .train(inputs.zip(sparseCat).filterNot(_._2 == null), rng = rng)
-      .getModel()
+      .train(sparseCatRows, rng = rng)
+      .model
       .transform(inputs)
-      .getExpected()
+      .expected
 
     val singleF1 = ClassificationMetrics.f1scores(reference, catLabel)
     val multiF1 = ClassificationMetrics.f1scores(catResults, catLabel)
@@ -98,16 +104,18 @@ class MultiTaskTreeTest extends SeedRandomMixIn {
     val combinedModelRng = Random(seed)
     val learner = MultiTaskTreeLearner()
     val combinedModelLearner = MultiTaskTreeLearner()
-    val models = learner.train(inputs.zip(labels), rng = trainRng).getModels()
-    val combinedModel = combinedModelLearner.train(inputs.zip(labels), rng = combinedModelRng).getModel()
+    val models = learner.train(TrainingRow.build(inputs.zip(labels)), rng = trainRng).models
+    val combinedModel =
+      combinedModelLearner.train(TrainingRow.build(inputs.zip(labels)), rng = combinedModelRng).model
 
     // Generate new inputs to test equality on.
-    val testInputs = TestUtils
-      .generateTrainingData(32, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
-      .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]]]
+    val testInputs = DataGenerator
+      .generate(32, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .data
+      .map(_.inputs)
+    val realResults = models.head.transform(testInputs).expected.asInstanceOf[Seq[Double]]
+    val catResults = models.last.transform(testInputs).expected.asInstanceOf[Seq[Boolean]]
+    val allResults = combinedModel.transform(testInputs).expected.asInstanceOf[Seq[Seq[Any]]]
     assert(Seq(realResults, catResults).transpose == allResults)
   }
 
@@ -115,9 +123,8 @@ class MultiTaskTreeTest extends SeedRandomMixIn {
   @Test
   def testFeatureImportance(): Unit = {
     val multiTaskLearner = MultiTaskTreeLearner()
-    val multiTaskTrainingResult = multiTaskLearner.train(inputs.zip(labels), rng = rng)
-    val importances = multiTaskTrainingResult.getFeatureImportance().get
+    val multiTaskTrainingResult = multiTaskLearner.train(TrainingRow.build(inputs.zip(labels)), rng = rng)
+    val importances = multiTaskTrainingResult.featureImportance.get
     assert(importances(1) == importances.max)
   }
-
 }
diff --git a/src/test/scala/io/citrine/lolo/trees/regression/RegressionTreeTest.scala b/src/test/scala/io/citrine/lolo/trees/regression/RegressionTreeTest.scala
index 9dca5b4e..481cc5f3 100644
--- a/src/test/scala/io/citrine/lolo/trees/regression/RegressionTreeTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/regression/RegressionTreeTest.scala
@@ -2,7 +2,8 @@ package io.citrine.lolo.trees.regression
 
 import java.io.{File, FileOutputStream, ObjectOutputStream}
 import breeze.linalg.DenseMatrix
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn, TestUtils}
 import io.citrine.lolo.linear.LinearRegressionLearner
 import io.citrine.lolo.stats.functions.Friedman
 import io.citrine.lolo.trees.splits.{BoltzmannSplitter, RegressionSplitter}
@@ -19,12 +20,12 @@ class RegressionTreeTest extends SeedRandomMixIn {
   def testFeatureImportanceNaN(): Unit = {
     val X = Vector.fill(100) {
       val input = Vector.fill(10)(1.0)
-      (input, 2.0)
+      TrainingRow(input, 2.0)
     }
 
     val DTLearner = RegressionTreeLearner()
     val DTMeta = DTLearner.train(X)
-    assert(DTMeta.getFeatureImportance().get.forall(v => !v.isNaN))
+    assert(DTMeta.featureImportance.get.forall(v => !v.isNaN))
   }
 
   /**
@@ -33,18 +34,18 @@ class RegressionTreeTest extends SeedRandomMixIn {
   @Test
   def testSimpleTree(): Unit = {
     val csv = TestUtils.readCsv("double_example.csv")
-    val trainingData = csv.map(vec => (vec.init, vec.last.asInstanceOf[Double]))
+    val trainingData = csv.map(vec => TrainingRow(vec.init, vec.last.asInstanceOf[Double]))
     val DTLearner = RegressionTreeLearner()
-    val DT = DTLearner.train(trainingData).getModel()
+    val DT = DTLearner.train(trainingData).model
 
     /* We should be able to memorize the inputs */
-    val output = DT.transform(trainingData.map(_._1))
-    trainingData.zip(output.getExpected()).foreach {
-      case ((_, a), p) =>
-        assert(Math.abs(a - p) < 1.0e-9)
+    val output = DT.transform(trainingData.map(_.inputs))
+    trainingData.zip(output.expected).foreach {
+      case (row, p) =>
+        assert(Math.abs(row.label - p) < 1.0e-9)
     }
-    assert(output.getGradient().isEmpty)
-    output.getDepth().foreach(d => assert(d > 3 && d < 9, s"Depth is $d"))
+    assert(output.gradient.isEmpty)
+    output.depth.foreach(d => assert(d > 3 && d < 9, s"Depth is $d"))
   }
 
   /**
@@ -54,18 +55,18 @@ class RegressionTreeTest extends SeedRandomMixIn {
   @Test
   def testrandomizePivotLocation(): Unit = {
     val csv = TestUtils.readCsv("double_example.csv")
-    val trainingData = csv.map(vec => (vec.init, vec.last.asInstanceOf[Double]))
+    val trainingData = csv.map(vec => TrainingRow(vec.init, vec.last.asInstanceOf[Double], 1.0))
     val DTLearner = RegressionTreeLearner(splitter = RegressionSplitter(randomizePivotLocation = true))
-    val DT = DTLearner.train(trainingData).getModel()
+    val DT = DTLearner.train(trainingData).model
 
     /* We should be able to memorize the inputs */
-    val output = DT.transform(trainingData.map(_._1))
-    trainingData.zip(output.getExpected()).foreach {
-      case ((x, a), p) =>
-        assert(Math.abs(a - p) < 1.0e-9)
+    val output = DT.transform(trainingData.map(_.inputs))
+    trainingData.zip(output.expected).foreach {
+      case (row, p) =>
+        assert(Math.abs(row.label - p) < 1.0e-9)
     }
-    assert(output.getGradient().isEmpty)
-    output.getDepth().foreach(d => assert(d > 3 && d < 9, s"Depth is $d"))
+    assert(output.gradient.isEmpty)
+    output.depth.foreach(d => assert(d > 3 && d < 9, s"Depth is $d"))
   }
 
   /**
@@ -74,28 +75,28 @@ class RegressionTreeTest extends SeedRandomMixIn {
   @Test
   def longerTest(): Unit = {
     val trainingData =
-      TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      DataGenerator.generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng).data
     val DTLearner = RegressionTreeLearner()
     val N = 100
     val start = System.nanoTime()
     val DTMeta = DTLearner.train(trainingData)
-    val DT = DTMeta.getModel()
-    (0 until N).map(i => DTLearner.train(trainingData).getModel())
+    val DT = DTMeta.model
+    (0 until N).map(i => DTLearner.train(trainingData).model)
     val duration = (System.nanoTime() - start) / 1.0e9
 
     println(s"Training large case took ${duration / N} s")
 
     /* We should be able to memorize the inputs */
-    val output = DT.transform(trainingData.map(_._1))
-    trainingData.zip(output.getExpected()).foreach {
-      case ((x, a), p) =>
-        assert(Math.abs(a - p) < 1.0e-9)
+    val output = DT.transform(trainingData.map(_.inputs))
+    trainingData.zip(output.expected).foreach {
+      case (row, p) =>
+        assert(Math.abs(row.label - p) < 1.0e-9)
     }
-    assert(output.getGradient().isEmpty)
-    output.getDepth().foreach(d => assert(d > 4 && d < 20, s"Depth is ${d}"))
+    assert(output.gradient.isEmpty)
+    output.depth.foreach(d => assert(d > 4 && d < 20, s"Depth is ${d}"))
 
     /* The first feature should be the most important */
-    val importances = DTMeta.getFeatureImportance().get
+    val importances = DTMeta.featureImportance.get
     assert(importances(1) == importances.max)
 
     val tmpFile: File = File.createTempFile("tmp", ".csv")
@@ -108,34 +109,32 @@ class RegressionTreeTest extends SeedRandomMixIn {
     */
   @Test
   def testCategorical(): Unit = {
-    val trainingData = TestUtils
-      .binTrainingData(
-        TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-        inputBins = Seq((0, 8))
-      )
-      .asInstanceOf[Seq[(Vector[Any], Double)]]
+    val trainingData = DataGenerator
+      .generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((0, 8)))
+      .data
 
     val DTLearner = RegressionTreeLearner()
     val N = 100
     val start = System.nanoTime()
     val DTMeta = DTLearner.train(trainingData)
-    val DT = DTMeta.getModel()
-    (0 until N).map(i => DTLearner.train(trainingData).getModel())
+    val DT = DTMeta.model
+    (0 until N).map(i => DTLearner.train(trainingData).model)
     val duration = (System.nanoTime() - start) / 1.0e9
 
     println(s"Training large case took ${duration / N} s")
 
     /* We should be able to memorize the inputs */
-    val output = DT.transform(trainingData.map(_._1))
-    trainingData.zip(output.getExpected()).foreach {
-      case ((x, a), p) =>
-        assert(Math.abs(a - p) < 1.0e-9)
+    val output = DT.transform(trainingData.map(_.inputs))
+    trainingData.zip(output.expected).foreach {
+      case (row, p) =>
+        assert(Math.abs(row.label - p) < 1.0e-9)
     }
-    assert(output.getGradient().isEmpty)
-    output.getDepth().foreach(d => assert(d > 4 && d < 21, s"Depth is ${d}"))
+    assert(output.gradient.isEmpty)
+    output.depth.foreach(d => assert(d > 4 && d < 21, s"Depth is ${d}"))
 
     /* The first feature should be the most important */
-    val importances = DTMeta.getFeatureImportance().get
+    val importances = DTMeta.featureImportance.get
     assert(importances(1) == importances.max)
     val tmpFile: File = File.createTempFile("tmp", ".csv")
     val oos = new ObjectOutputStream(new FileOutputStream(tmpFile))
@@ -148,24 +147,24 @@ class RegressionTreeTest extends SeedRandomMixIn {
   @Test
   def testLinearLeaves(): Unit = {
     val trainingData =
-      TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      DataGenerator.generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng).data
 
     val linearLearner = LinearRegressionLearner(regParam = Some(0.0))
     val DTLearner = RegressionTreeLearner(leafLearner = Some(linearLearner), minLeafInstances = 2)
     val DTMeta = DTLearner.train(trainingData)
-    val DT = DTMeta.getModel()
+    val DT = DTMeta.model
 
     /* We should be able to memorize the inputs */
-    val output = DT.transform(trainingData.map(_._1))
-    trainingData.zip(output.getExpected()).foreach {
-      case ((x, a), p) =>
-        assert(Math.abs(a - p) < 1.0e-9)
+    val output = DT.transform(trainingData.map(_.inputs))
+    trainingData.zip(output.expected).foreach {
+      case (row, p) =>
+        assert(Math.abs(row.label - p) < 1.0e-9)
     }
-    assert(output.getGradient().isDefined)
-    output.getDepth().foreach(d => assert(d > 4 && d < 18, s"Depth is $d"))
+    assert(output.gradient.isDefined)
+    output.depth.foreach(d => assert(d > 4 && d < 18, s"Depth is $d"))
 
     /* The first feature should be the most important */
-    val importances = DTMeta.getFeatureImportance().get
+    val importances = DTMeta.featureImportance.get
     assert(importances(1) == importances.max)
     /* They should all be non-zero */
     assert(importances.min > 0.0)
@@ -180,21 +179,19 @@ class RegressionTreeTest extends SeedRandomMixIn {
     */
   @Test
   def testStumpWithLinearLeaf(): Unit = {
-    val trainingData = TestUtils
-      .binTrainingData(
-        TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng),
-        inputBins = Seq((11, 8))
-      )
-      .asInstanceOf[Seq[(Vector[Any], Double)]]
+    val trainingData = DataGenerator
+      .generate(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
+      .withBinnedInputs(bins = Seq((11, 8)))
+      .data
 
     val linearLearner = LinearRegressionLearner(regParam = Some(1.0))
     val DTLearner = RegressionTreeLearner(leafLearner = Some(linearLearner), maxDepth = 0)
     val DTMeta = DTLearner.train(trainingData)
-    val DT = DTMeta.getModel()
+    val DT = DTMeta.model
 
-    val linearImportance = linearLearner.train(trainingData).getFeatureImportance().get
+    val linearImportance = linearLearner.train(trainingData).featureImportance.get
 
-    val importances = DTMeta.getFeatureImportance().get
+    val importances = DTMeta.featureImportance.get
 
     /* The first feature should be the most important */
     assert(importances(1) == importances.max)
@@ -206,8 +203,8 @@ class RegressionTreeTest extends SeedRandomMixIn {
       s"Expected linear and maxDepth=0 importances to align"
     )
 
-    val result = DT.transform(trainingData.map(_._1))
-    assert(result.getDepth().forall(_ == 0), s"Expected all the predictions to be depth 0")
+    val result = DT.transform(trainingData.map(_.inputs))
+    assert(result.depth.forall(_ == 0), s"Expected all the predictions to be depth 0")
   }
 
   /**
@@ -216,20 +213,15 @@ class RegressionTreeTest extends SeedRandomMixIn {
   @Test
   def testWeights(): Unit = {
     val trainingData =
-      TestUtils.generateTrainingData(32, 12, noise = 100.0, function = Friedman.friedmanSilverman, rng = rng)
+      DataGenerator.generate(32, 12, noise = 100.0, function = Friedman.friedmanSilverman, rng = rng).data
+    val weightedData = trainingData.map(_.withWeight(rng.nextInt(8).toDouble))
 
     val linearLearner = LinearRegressionLearner(regParam = Some(1.0))
     val DTLearner = RegressionTreeLearner(leafLearner = Some(linearLearner), maxDepth = 1)
-    val DTMeta = DTLearner.train(
-      trainingData,
-      weights = Some(Seq.fill(trainingData.size) {
-        rng.nextInt(8)
-      })
-    )
-    val DT = DTMeta.getModel()
+    val DTMeta = DTLearner.train(weightedData)
 
     /* The first feature should be the most important */
-    val importances = DTMeta.getFeatureImportance().get
+    val importances = DTMeta.featureImportance.get
     assert(importances.forall(_ >= 0.0), "Found negative feature importance")
   }
 
@@ -239,17 +231,17 @@ class RegressionTreeTest extends SeedRandomMixIn {
   @Test
   def testSimpleBoltzmannTree(): Unit = {
     val csv = TestUtils.readCsv("double_example.csv")
-    val trainingData = csv.map(vec => (vec.init, vec.last.asInstanceOf[Double]))
+    val trainingData = csv.map(vec => TrainingRow(vec.init, vec.last.asInstanceOf[Double], 1.0))
     val DTLearner = RegressionTreeLearner(splitter = BoltzmannSplitter(1e-4))
-    val DT = DTLearner.train(trainingData, rng = rng).getModel()
+    val DT = DTLearner.train(trainingData, rng = rng).model
 
     /* We should be able to memorize the inputs */
-    val output = DT.transform(trainingData.map(_._1))
-    trainingData.zip(output.getExpected()).foreach {
-      case ((_, a), p) => assert(Math.abs(a - p) < 1.0e-9)
+    val output = DT.transform(trainingData.map(_.inputs))
+    trainingData.zip(output.expected).foreach {
+      case (row, p) => assert(Math.abs(row.label - p) < 1.0e-9)
     }
-    assert(output.getGradient().isEmpty)
-    output.getDepth().zip(output.getExpected()).foreach {
+    assert(output.gradient.isEmpty)
+    output.depth.zip(output.expected).foreach {
       case (d, y) => assert(d > 3 && d < 9, s"Depth is $d at y=$y")
     }
   }
@@ -262,16 +254,16 @@ class RegressionTreeTest extends SeedRandomMixIn {
     * @param expected     known reference Shapley value to which to compare (fails upon mismatch)
     */
   def shapleyCompare(
-      trainingData: Seq[(Vector[Double], Double)],
+      trainingData: Seq[TrainingRow[Double]],
       evalLocation: Vector[Any],
       expected: Vector[Double],
       omitFeatures: Set[Int] = Set()
   ): Unit = {
-    val actual = RegressionTreeLearner().train(trainingData).getModel().shapley(evalLocation, omitFeatures) match {
+    val actual = RegressionTreeLearner().train(trainingData).model.shapley(evalLocation, omitFeatures) match {
       case None => fail("Unexpected None returned by shapley.")
       case x: Option[DenseMatrix[Double]] => {
         val a = x.get
-        assert(a.cols == trainingData.head._1.length, "Expected one Shapley value per feature.")
+        assert(a.cols == trainingData.head.inputs.length, "Expected one Shapley value per feature.")
         assert(a.rows == 1, "Expected a single output dimension.")
         a.toDenseVector.toScalaVector
       }
@@ -292,20 +284,20 @@ class RegressionTreeTest extends SeedRandomMixIn {
   def testShapley(): Unit = {
     // Example from Lundberg paper (https://arxiv.org/pdf/1802.03888.pdf)
     val trainingData1 = Seq(
-      (Vector(1.0, 1.0), 80.0),
-      (Vector(1.0, 0.0), 0.0),
-      (Vector(0.0, 1.0), 0.0),
-      (Vector(0.0, 0.0), 0.0)
+      TrainingRow(Vector(1.0, 1.0), 80.0),
+      TrainingRow(Vector(1.0, 0.0), 0.0),
+      TrainingRow(Vector(0.0, 1.0), 0.0),
+      TrainingRow(Vector(0.0, 0.0), 0.0)
     )
     val expected1 = Vector(30.0, 30.0)
     shapleyCompare(trainingData1, Vector[Any](1.0, 1.0), expected1)
 
     // Second example from Lundberg paper (https://arxiv.org/pdf/1802.03888.pdf)
     val trainingData2 = Seq(
-      (Vector(1.0, 1.0), 90.0),
-      (Vector(1.0, 0.0), 10.0),
-      (Vector(0.0, 1.0), 0.0),
-      (Vector(0.0, 0.0), 0.0)
+      TrainingRow(Vector(1.0, 1.0), 90.0),
+      TrainingRow(Vector(1.0, 0.0), 10.0),
+      TrainingRow(Vector(0.0, 1.0), 0.0),
+      TrainingRow(Vector(0.0, 0.0), 0.0)
     )
     val expected2 = Vector(35.0, 30.0)
     shapleyCompare(trainingData2, Vector[Any](1.0, 1.0), expected2)
@@ -313,12 +305,12 @@ class RegressionTreeTest extends SeedRandomMixIn {
     // Example with two splits on one feature
     // Worked out with pen-and-paper from Lundberg Equation 2.
     val trainingData3 = Seq(
-      (Vector(1.0, 1.0), 100.0),
-      (Vector(1.0, 0.0), 80.0),
-      (Vector(1.0, 0.2), 70.0),
-      (Vector(0.0, 1.0), 0.0),
-      (Vector(0.0, 0.2), 0.0),
-      (Vector(0.0, 0.0), 0.0)
+      TrainingRow(Vector(1.0, 1.0), 100.0),
+      TrainingRow(Vector(1.0, 0.0), 80.0),
+      TrainingRow(Vector(1.0, 0.2), 70.0),
+      TrainingRow(Vector(0.0, 1.0), 0.0),
+      TrainingRow(Vector(0.0, 0.2), 0.0),
+      TrainingRow(Vector(0.0, 0.0), 0.0)
     )
     val expected3 = Vector(45.8333333333333, 12.5)
     shapleyCompare(trainingData3, Vector[Any](1.0, 1.0), expected3)
@@ -326,12 +318,12 @@ class RegressionTreeTest extends SeedRandomMixIn {
     // Example with 5 features, to exercise all the factorials in Lundberg equation 2.
     // Referenced against the shap package on a sklearn decision tree.
     val trainingData4 = Seq(
-      (Vector(0.0, 0.0, 0.0, 0.0, 0.0), 1.0),
-      (Vector(1.0, 0.0, 0.0, 0.0, 0.0), 2.0),
-      (Vector(0.0, 1.0, 0.0, 0.0, 0.0), 4.0),
-      (Vector(0.0, 0.0, 1.0, 0.0, 0.0), 8.0),
-      (Vector(0.0, 0.0, 0.0, 1.0, 0.0), 16.0),
-      (Vector(0.0, 0.0, 0.0, 0.0, 1.0), 32.0)
+      TrainingRow(Vector(0.0, 0.0, 0.0, 0.0, 0.0), 1.0),
+      TrainingRow(Vector(1.0, 0.0, 0.0, 0.0, 0.0), 2.0),
+      TrainingRow(Vector(0.0, 1.0, 0.0, 0.0, 0.0), 4.0),
+      TrainingRow(Vector(0.0, 0.0, 1.0, 0.0, 0.0), 8.0),
+      TrainingRow(Vector(0.0, 0.0, 0.0, 1.0, 0.0), 16.0),
+      TrainingRow(Vector(0.0, 0.0, 0.0, 0.0, 1.0), 32.0)
     )
     val expected4 = Vector(0.0333333333333333, 0.2, 0.8666666666666667, 3.533333333333333, 16.866666666666667)
     shapleyCompare(trainingData4, Vector.fill[Any](5)(1.0), expected4)
@@ -343,6 +335,6 @@ class RegressionTreeTest extends SeedRandomMixIn {
     shapleyCompare(trainingData2, Vector[Any](1.0, 1.0), expected2b, omitFeatures = Set(1))
 
     // Ensure we don't crash when restricting number of features.
-    RegressionTreeLearner(numFeatures = 1).train(trainingData4).getModel().shapley(Vector.fill[Any](5)(0.0), Set())
+    RegressionTreeLearner(numFeatures = 1).train(trainingData4).model.shapley(Vector.fill[Any](5)(0.0), Set())
   }
 }
diff --git a/src/test/scala/io/citrine/lolo/trees/splits/BoltzmannSplitterTest.scala b/src/test/scala/io/citrine/lolo/trees/splits/BoltzmannSplitterTest.scala
index f6468539..97b86f65 100644
--- a/src/test/scala/io/citrine/lolo/trees/splits/BoltzmannSplitterTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/splits/BoltzmannSplitterTest.scala
@@ -1,6 +1,7 @@
 package io.citrine.lolo.trees.splits
 
 import io.citrine.lolo.SeedRandomMixIn
+import io.citrine.lolo.api.TrainingRow
 import io.citrine.random.Random
 import org.junit.Test
 
@@ -16,7 +17,7 @@ class BoltzmannSplitterTest extends SeedRandomMixIn {
       val x = rng.nextDouble()
       val y = 1.0
       val weight = 1.0
-      (Vector(x), y, weight)
+      TrainingRow(Vector(x), y, weight)
     }
 
     val (bestSplit, improvement) = splitter.getBestSplit(testData, 1, 1)
@@ -37,7 +38,7 @@ class BoltzmannSplitterTest extends SeedRandomMixIn {
       val x = rng.nextDouble()
       val y = rng.nextGaussian() * 1.0e-9 + 1.0
       val weight = 1.0
-      (Vector(x), y, weight)
+      TrainingRow(Vector(x), y, weight)
     }
 
     splitter.getBestSplit(testData, 1, 1, rng = rng)
diff --git a/src/test/scala/io/citrine/lolo/trees/splits/ClassificationSplitterTest.scala b/src/test/scala/io/citrine/lolo/trees/splits/ClassificationSplitterTest.scala
index a092da02..a220f6c2 100644
--- a/src/test/scala/io/citrine/lolo/trees/splits/ClassificationSplitterTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/splits/ClassificationSplitterTest.scala
@@ -1,6 +1,7 @@
 package io.citrine.lolo.trees.splits
 
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
 import io.citrine.lolo.encoders.CategoricalEncoder
 import io.citrine.theta.Stopwatch
 
@@ -16,18 +17,13 @@ class ClassificationSplitterTest extends SeedRandomMixIn {
   val nLabel = 4096
   val nSubset = 4
 
-  val testData: Seq[(Vector[Any], Any)] = TestUtils.binTrainingData(
-    TestUtils.generateTrainingData(nRow, 12, rng = rng),
-    responseBins = Some(nLabel)
-  )
+  val testData: Seq[TrainingRow[String]] = DataGenerator.generate(nRow, 12, rng = rng).withBinnedLabels(nLabel).data
 
-  val encoder: CategoricalEncoder[Any] = CategoricalEncoder.buildEncoder(testData.map(_._2))
-  val encodedData: Seq[(Vector[AnyVal], Char, Double)] = testData.map {
-    case (f, l) =>
-      (f.asInstanceOf[Vector[AnyVal]], encoder.encode(l), 1.0)
+  val encoder: CategoricalEncoder[Any] = CategoricalEncoder.buildEncoder(testData.map(_.label))
+  val encodedData: Seq[TrainingRow[Char]] = testData.map { row =>
+    row.withLabel(encoder.encode(row.getClass))
   }
-
-  val subset: Seq[(Vector[AnyVal], Char, Double)] = encodedData.take(nSubset)
+  val subset: Seq[TrainingRow[Char]] = encodedData.take(nSubset)
 
   /**
     * Evaluate the split finding performance on large and small datasets
diff --git a/src/test/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitterTest.scala b/src/test/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitterTest.scala
index e80ea270..95e95285 100644
--- a/src/test/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitterTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitterTest.scala
@@ -1,6 +1,7 @@
 package io.citrine.lolo.trees.splits
 
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
+import io.citrine.lolo.api.TrainingRow
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
 import io.citrine.random.Random
 import org.junit.Test
 
@@ -16,7 +17,7 @@ class ExtraRandomRegressionSplitterTest extends SeedRandomMixIn {
       val x = rng.nextDouble()
       val y = 1.0
       val weight = 1.0
-      (Vector(x), y, weight)
+      TrainingRow(Vector(x), y, weight)
     }
 
     val (bestSplit, improvement) = splitter.getBestSplit(testData, 1, 1)
@@ -37,7 +38,7 @@ class ExtraRandomRegressionSplitterTest extends SeedRandomMixIn {
       val x = rng.nextDouble()
       val y = rng.nextGaussian() * 1.0e-9 + 1.0
       val weight = 1.0
-      (Vector(x), y, weight)
+      TrainingRow(Vector(x), y, weight)
     }
 
     splitter.getBestSplit(testData, 1, 1, rng = rng)
@@ -54,7 +55,7 @@ class ExtraRandomRegressionSplitterTest extends SeedRandomMixIn {
     Seq(1, 5, 9).foreach { numFeatures =>
       // Arrange training input data as a regular grid, and then as a Gaussian-distributed sample.
       Seq(
-        TestUtils.enumerateGrid(Seq.fill(numFeatures)(baseGrid)),
+        DataGenerator.enumerateGrid(Seq.fill(numFeatures)(baseGrid)),
         Seq.fill(64)(Vector.fill(numFeatures)(rng.nextGaussian()))
       ).foreach { xTrain =>
         // Have ExtraRandomSplitter.getBestSplit choose a split from this number of randomly-chosen features.
@@ -67,7 +68,7 @@ class ExtraRandomRegressionSplitterTest extends SeedRandomMixIn {
 
             // Generate linear training data, using xTrain from above as inputs, starting by selecting coefficients for each feature.
             val featureCoefficients = rng.shuffle(featureIndices).map { i => Math.pow((i + 1).toDouble, 2.0) }
-            val trainingData = xTrain.map { x =>
+            val trainingTuples = xTrain.map { x =>
               val y = x.zip(featureCoefficients).map { case (a, b) => a * b }.sum
               val weight = 1.0
               (x, y, weight)
@@ -86,7 +87,7 @@ class ExtraRandomRegressionSplitterTest extends SeedRandomMixIn {
               .zip(randomUniforms)
               .map {
                 case (i, u) =>
-                  val x = trainingData.map(_._1(i))
+                  val x = trainingTuples.map(_._1(i))
                   val xmin = x.min
                   val xmax = x.max
                   val cutPoint = xmin + u * (xmax - xmin)
@@ -97,7 +98,7 @@ class ExtraRandomRegressionSplitterTest extends SeedRandomMixIn {
 
             // Compute the sum of variances across partitions of each possible cut.
             val varianceSums = featureIndices.map { k =>
-              trainingData
+              trainingTuples
                 .groupBy { v => v._1(k) < cutPoints(k) }
                 .flatMap {
                   case (_, subset) =>
@@ -123,7 +124,8 @@ class ExtraRandomRegressionSplitterTest extends SeedRandomMixIn {
             rng = Random(sharedSeed)
             val splitter = ExtraRandomRegressionSplitter()
             // Ask the splitter what it chose as a split.
-            val bestSplit = splitter.getBestSplit(trainingData, numFeaturesToConsider, 1, rng = rng)
+            val bestSplit =
+              splitter.getBestSplit(TrainingRow.build(trainingTuples), numFeaturesToConsider, 1, rng = rng)
 
             // Finally, we can ensure that the index on which we split is correct...
             val testCaveatMessage = "NOTE: this test may inaccurately fail due to changes in the sequence of rng calls."
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 04ceb67e..1219eafe 100644
--- a/src/test/scala/io/citrine/lolo/trees/splits/MultiTaskSplitterTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/splits/MultiTaskSplitterTest.scala
@@ -1,6 +1,7 @@
 package io.citrine.lolo.trees.splits
 
 import io.citrine.lolo.SeedRandomMixIn
+import io.citrine.lolo.api.TrainingRow
 import io.citrine.lolo.trees.impurity.MultiImpurityCalculator
 import org.junit.Test
 
@@ -16,16 +17,15 @@ class MultiTaskSplitterTest extends SeedRandomMixIn {
   @Test
   def testBestRealSplit(): Unit = {
     val data = Seq(
-      (Vector(0.0), Vector[AnyVal](0.0, 1.toChar), 1.0),
-      (Vector(1.0), Vector[AnyVal](1.0, 2.toChar), 1.0),
-      (Vector(2.0), Vector[AnyVal](2.0, 2.toChar), 1.0),
-      (Vector(100.0), Vector[AnyVal](200.0, 3.toChar), 0.0)
+      TrainingRow(Vector(0.0), Vector[AnyVal](0.0, 1.toChar), 1.0),
+      TrainingRow(Vector(1.0), Vector[AnyVal](1.0, 2.toChar), 1.0),
+      TrainingRow(Vector(2.0), Vector[AnyVal](2.0, 2.toChar), 1.0),
+      TrainingRow(Vector(100.0), Vector[AnyVal](200.0, 3.toChar), 0.0)
     )
-    val calculator = MultiImpurityCalculator.build(data.map(_._2), data.map(_._3))
+    val calculator = MultiImpurityCalculator.build(data.map(_.label), data.map(_.weight))
 
-    val splitter = MultiTaskSplitter()
     val (pivot, impurity) = Splitter.getBestRealSplit(data, calculator, 0, 1, rng = rng)
-    val turns = data.map(x => pivot.turnLeft(x._1))
+    val turns = data.map(x => pivot.turnLeft(x.inputs))
     assert(turns == Seq(true, false, false, false))
     assert(impurity == 0.5)
   }
@@ -36,16 +36,16 @@ class MultiTaskSplitterTest extends SeedRandomMixIn {
   @Test
   def testBestCategoricalSplit(): Unit = {
     val data = Seq(
-      (Vector(1.toChar), Vector[AnyVal](0.0, 1.toChar), 1.0),
-      (Vector(2.toChar), Vector[AnyVal](1.0, 2.toChar), 1.0),
-      (Vector(2.toChar), Vector[AnyVal](2.0, 2.toChar), 1.0),
-      (Vector(3.toChar), Vector[AnyVal](200.0, 3.toChar), 0.0)
+      TrainingRow(Vector(1.toChar), Vector[AnyVal](0.0, 1.toChar), 1.0),
+      TrainingRow(Vector(2.toChar), Vector[AnyVal](1.0, 2.toChar), 1.0),
+      TrainingRow(Vector(2.toChar), Vector[AnyVal](2.0, 2.toChar), 1.0),
+      TrainingRow(Vector(3.toChar), Vector[AnyVal](200.0, 3.toChar), 0.0)
     )
-    val calculator = MultiImpurityCalculator.build(data.map(_._2), data.map(_._3))
+    val calculator = MultiImpurityCalculator.build(data.map(_.label), data.map(_.weight))
 
     val splitter = MultiTaskSplitter()
     val (pivot, impurity) = splitter.getBestCategoricalSplit(data, calculator, 0, 1)
-    val turns = data.map(x => pivot.turnLeft(x._1)).take(3)
+    val turns = data.map(x => pivot.turnLeft(x.inputs)).take(3)
     println(turns)
     assert(turns == Seq(true, false, false))
     assert(impurity == 0.5)
@@ -63,10 +63,10 @@ class MultiTaskSplitterTest extends SeedRandomMixIn {
     val labels = Vector(Vector(1.238, null), Vector(1.180, null))
     val weights = Vector(1.0, 1.0)
     val data = inputs.indices.map { i =>
-      (inputs(i), labels(i).asInstanceOf[Vector[AnyVal]], weights(i))
+      TrainingRow(inputs(i), labels(i).asInstanceOf[Vector[AnyVal]], weights(i))
     }
     val splitter = MultiTaskSplitter()
-    val (pivot, _) = splitter.getBestSplit(data, data.head._1.size, 1, rng = rng)
+    val (pivot, _) = splitter.getBestSplit(data, data.head.inputs.size, 1, rng = rng)
     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 55bc899a..c6578c31 100644
--- a/src/test/scala/io/citrine/lolo/trees/splits/SplitterTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/splits/SplitterTest.scala
@@ -1,12 +1,10 @@
 package io.citrine.lolo.trees.splits
 
 import io.citrine.lolo.SeedRandomMixIn
+import io.citrine.lolo.api.TrainingRow
 import io.citrine.lolo.trees.impurity.VarianceCalculator
 import org.junit.Test
 
-/**
-  * Created by maxhutch on 12/1/16.
-  */
 @Test
 class SplitterTest extends SeedRandomMixIn {
 
@@ -59,7 +57,7 @@ class SplitterTest extends SeedRandomMixIn {
   def testLargeDuplicates(): Unit = {
     val base: Double = 3.0e9
     val trainingData = Seq.fill(8) {
-      (Vector(base + rng.nextDouble()), rng.nextDouble(), 1.0)
+      TrainingRow(Vector(base + rng.nextDouble()), rng.nextDouble())
     }
 
     val calculator = new VarianceCalculator(0.0, 0.0, 8.0)
diff --git a/src/test/scala/io/citrine/lolo/util/LoloPyDataLoaderTest.scala b/src/test/scala/io/citrine/lolo/util/LoloPyDataLoaderTest.scala
index 9823feae..f4ee0fae 100644
--- a/src/test/scala/io/citrine/lolo/util/LoloPyDataLoaderTest.scala
+++ b/src/test/scala/io/citrine/lolo/util/LoloPyDataLoaderTest.scala
@@ -1,7 +1,7 @@
 package io.citrine.lolo.util
-import io.citrine.lolo.{PredictionResult, TestUtils}
-import io.citrine.lolo.learners.RandomForest
-import io.citrine.lolo.stats.functions.Friedman
+
+import io.citrine.lolo.api.PredictionResult
+import io.citrine.lolo.learners.RandomForestRegressor
 import org.junit.Test
 
 class LoloPyDataLoaderTest {
@@ -40,18 +40,18 @@ class LoloPyDataLoaderTest {
   @Test
   def testGetRegression(): Unit = {
     val results = new PredictionResult[Double] {
-      override def getExpected(): Seq[Double] = {
+      override def expected: Seq[Double] = {
         1.0 :: 2.0 :: 3.0 :: Nil
       }
 
-      override def getUncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = {
+      override def uncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = {
         Some(0.1 :: 0.2 :: 0.3 :: Nil)
       }
     }
 
     // Get the expected results
-    val means = results.getExpected()
-    val sigma = results.getUncertainty().get.asInstanceOf[Seq[Double]]
+    val means = results.expected
+    val sigma = results.uncertainty().get.asInstanceOf[Seq[Double]]
 
     // Get the results as a byte array, and convert them back
     val reproMeans = LoloPyDataLoader
@@ -70,21 +70,21 @@ class LoloPyDataLoaderTest {
   def testGetClassification(): Unit = {
     // Get the base results
     val results = new PredictionResult[Int] {
-      override def getExpected(): Seq[Int] = {
+      override def expected: Seq[Int] = {
         0 :: 1 :: 0 :: 1 :: Nil
       }
 
-      override def getUncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = {
+      override def uncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = {
         Some(Map(0 -> 1.0) :: Map(1 -> 0.5, 0 -> 0.2, 2 -> 0.3) :: Map(0 -> 1.0) :: Map(1 -> 1.0) :: Nil)
       }
 
-      override def getImportanceScores(): Option[Seq[Seq[Double]]] = {
+      override def importanceScores: Option[Seq[Seq[Double]]] = {
         Some((0.1 :: 0.2 :: Nil) :: (0.4 :: 0.8 :: Nil) :: (0.1 :: 0.5 :: Nil) :: (0.8 :: 0.2 :: Nil) :: Nil)
       }
     }
-    val means = results.getExpected()
+    val means = results.expected
     val probs: Seq[Double] = results
-      .getUncertainty()
+      .uncertainty()
       .get
       .asInstanceOf[Seq[Map[Int, Double]]]
       .flatMap(x => 0.until(3).map(i => x.getOrElse(i, 0.0)))
@@ -119,9 +119,9 @@ class LoloPyDataLoaderTest {
 
   @Test
   def testSerialization(): Unit = {
-    val model = new RandomForest(numTrees = 256)
+    val model = new RandomForestRegressor(numTrees = 256)
     val bytes = LoloPyDataLoader.serializeObject(model, 9)
-    val model2 = LoloPyDataLoader.deserializeObject(bytes).asInstanceOf[RandomForest]
+    val model2 = LoloPyDataLoader.deserializeObject(bytes).asInstanceOf[RandomForestRegressor]
     assert(model.numTrees == model2.numTrees)
   }
 }
diff --git a/src/test/scala/io/citrine/lolo/validation/CalibrationStudy.scala b/src/test/scala/io/citrine/lolo/validation/CalibrationStudy.scala
index 313c6337..90723e53 100644
--- a/src/test/scala/io/citrine/lolo/validation/CalibrationStudy.scala
+++ b/src/test/scala/io/citrine/lolo/validation/CalibrationStudy.scala
@@ -1,9 +1,10 @@
 package io.citrine.lolo.validation
 
-import io.citrine.lolo.bags.Bagger
+import io.citrine.lolo.api.{PredictionResult, TrainingRow}
+import io.citrine.lolo.bags.{Bagger, RegressionBagger}
 import io.citrine.lolo.stats.functions.{Friedman, Linear}
 import io.citrine.lolo.trees.regression.RegressionTreeLearner
-import io.citrine.lolo.{PredictionResult, SeedRandomMixIn, TestUtils}
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn, TestUtils}
 import org.apache.commons.math3.distribution.CauchyDistribution
 import org.knowm.xchart.BitmapEncoder.BitmapFormat
 import org.knowm.xchart.{BitmapEncoder, CategoryChart, CategoryChartBuilder}
@@ -54,7 +55,7 @@ object CalibrationStudy extends SeedRandomMixIn {
       ratios: Seq[Int] = Seq(1, 2, 4)
   ): Unit = {
     val nFeature = 8
-    val data = TestUtils.iterateTrainingData(nFeature, Friedman.friedmanSilverman, rng = rng)
+    val data = DataGenerator.iterate(nFeature, Friedman.friedmanSilverman, rng = rng)
 
     ratios.foreach { ratio =>
       val chart = Merit.plotMeritScan(
@@ -72,10 +73,8 @@ object CalibrationStudy extends SeedRandomMixIn {
         rng = rng
       ) { nTrain: Double =>
         val nTree = ratio * nTrain
-        val learner = Bagger(
-          RegressionTreeLearner(
-            numFeatures = nFeature
-          ),
+        val learner = RegressionBagger(
+          RegressionTreeLearner(numFeatures = nFeature),
           numBags = nTree.toInt,
           useJackknife = true,
           uncertaintyCalibration = false
@@ -92,7 +91,7 @@ object CalibrationStudy extends SeedRandomMixIn {
       calibrated: Boolean = false
   ): Unit = {
     val nFeature = 8
-    val data = TestUtils.iterateTrainingData(nFeature, Friedman.friedmanSilverman, rng = rng)
+    val data = DataGenerator.iterate(nFeature, Friedman.friedmanSilverman, rng = rng)
     sizes.foreach { nTrain =>
       val chart = Merit.plotMeritScan(
         "Number of trees",
@@ -108,10 +107,8 @@ object CalibrationStudy extends SeedRandomMixIn {
         yMax = Some(1.0),
         rng = rng
       ) { nTree: Double =>
-        val learner = Bagger(
-          RegressionTreeLearner(
-            numFeatures = nFeature
-          ),
+        val learner = RegressionBagger(
+          RegressionTreeLearner(numFeatures = nFeature),
           numBags = nTree.toInt,
           useJackknife = true,
           uncertaintyCalibration = calibrated
@@ -135,7 +132,7 @@ object CalibrationStudy extends SeedRandomMixIn {
       funcName: String = "fs"
   ): Unit = {
     val nFeature = 8
-    val data = TestUtils.iterateTrainingData(nFeature, func, rng = rng)
+    val data = DataGenerator.iterate(nFeature, func, rng = rng)
     sizes.foreach { nTree =>
       val chart = Merit.plotMeritScan(
         "Number of training points",
@@ -151,10 +148,8 @@ object CalibrationStudy extends SeedRandomMixIn {
         yMax = Some(1.0),
         rng = rng
       ) { nTrain: Double =>
-        val learner = Bagger(
-          RegressionTreeLearner(
-            numFeatures = nFeature
-          ),
+        val learner = RegressionBagger(
+          RegressionTreeLearner(numFeatures = nFeature),
           numBags = nTree.toInt,
           useJackknife = true,
           uncertaintyCalibration = calibrated
@@ -181,13 +176,11 @@ object CalibrationStudy extends SeedRandomMixIn {
       ignoreDims: Int = 0
   ): Map[String, (Double, Double)] = {
 
-    val data = TestUtils
-      .iterateTrainingData(nFeature, func, rng = rng)
-      .map { case (x, y) => (x.drop(ignoreDims), y) }
-    val learner = Bagger(
-      RegressionTreeLearner(
-        numFeatures = nFeature
-      ),
+    val data = DataGenerator
+      .iterate(nFeature, func, rng = rng)
+      .map { row => row.mapInputs(_.drop(ignoreDims)) }
+    val learner = RegressionBagger(
+      RegressionTreeLearner(numFeatures = nFeature),
       numBags = nTree,
       useJackknife = true,
       uncertaintyCalibration = true,
@@ -244,9 +237,8 @@ object CalibrationStudy extends SeedRandomMixIn {
   ): CategoryChart = {
     val data: Seq[(Double, Double, Double)] = source.flatMap {
       case (predictions, actual) =>
-        predictions
-          .getExpected()
-          .lazyZip(predictions.getUncertainty().get.asInstanceOf[Seq[Double]])
+        predictions.expected
+          .lazyZip(predictions.uncertainty().get.asInstanceOf[Seq[Double]])
           .lazyZip(actual)
           .toSeq
     }.toSeq
@@ -302,7 +294,7 @@ object CalibrationStudy extends SeedRandomMixIn {
       calibrated: Boolean = false
   ): Unit = {
     val csv = TestUtils.readCsv("hcep.csv")
-    val trainingData = csv.tail.map(vec => (vec.init, vec.last.asInstanceOf[Double]))
+    val trainingData = csv.tail.map(vec => TrainingRow(vec.init, vec.last.asInstanceOf[Double]))
     val nFeature = 8
     sizes.foreach { nTree =>
       val chart = Merit.plotMeritScan(
@@ -319,10 +311,8 @@ object CalibrationStudy extends SeedRandomMixIn {
         yMax = Some(1.0),
         rng = rng
       ) { nTrain: Double =>
-        val learner = Bagger(
-          RegressionTreeLearner(
-            numFeatures = nFeature
-          ),
+        val learner = RegressionBagger(
+          RegressionTreeLearner(numFeatures = nFeature),
           numBags = nTree.toInt,
           useJackknife = true,
           uncertaintyCalibration = calibrated
@@ -354,12 +344,10 @@ object CalibrationStudy extends SeedRandomMixIn {
   ): Map[String, (Double, Double)] = {
 
     val csv = TestUtils.readCsv("hcep.csv")
-    val trainingData = csv.tail.map(vec => (vec.init, vec.last.asInstanceOf[Double]))
+    val trainingData = csv.tail.map(vec => TrainingRow(vec.init, vec.last.asInstanceOf[Double]))
 
-    val learner = Bagger(
-      RegressionTreeLearner(
-        numFeatures = nFeature / 3
-      ),
+    val learner = RegressionBagger(
+      RegressionTreeLearner(numFeatures = nFeature / 3),
       numBags = nTree,
       useJackknife = true,
       uncertaintyCalibration = false
diff --git a/src/test/scala/io/citrine/lolo/validation/CrossValidationTest.scala b/src/test/scala/io/citrine/lolo/validation/CrossValidationTest.scala
index 9e194434..39c2bc6f 100644
--- a/src/test/scala/io/citrine/lolo/validation/CrossValidationTest.scala
+++ b/src/test/scala/io/citrine/lolo/validation/CrossValidationTest.scala
@@ -1,7 +1,7 @@
 package io.citrine.lolo.validation
 
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
-import io.citrine.lolo.learners.RandomForest
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
+import io.citrine.lolo.learners.RandomForestRegressor
 import io.citrine.lolo.stats.functions.Friedman
 import org.junit.Test
 
@@ -14,8 +14,8 @@ class CrossValidationTest extends SeedRandomMixIn {
     */
   @Test
   def testCompareToOutOfBag(): Unit = {
-    val learner = RandomForest()
-    val data = TestUtils.generateTrainingData(128, 8, Friedman.friedmanSilverman, rng = rng)
+    val learner = RandomForestRegressor()
+    val data = DataGenerator.generate(128, 8, Friedman.friedmanSilverman, rng = rng).data
 
     val metric = RootMeanSquareError
     val (rmseFromCV, uncertainty) =
@@ -23,13 +23,9 @@ class CrossValidationTest extends SeedRandomMixIn {
 
     val trainingResult = learner.train(data)
     val rmseFromPVA = Math.sqrt(
-      trainingResult
-        .getPredictedVsActual()
-        .get
-        .map {
-          case (_, p: Double, a: Double) => Math.pow(p - a, 2.0)
-        }
-        .sum / trainingResult.getPredictedVsActual().get.size
+      trainingResult.predictedVsActual.get.map {
+        case (_, p: Double, a: Double) => Math.pow(p - a, 2.0)
+      }.sum / trainingResult.predictedVsActual.get.size
     )
 
     // These have a false negative rate less than 1/100 at the time of original authorship
diff --git a/src/test/scala/io/citrine/lolo/validation/MeritTest.scala b/src/test/scala/io/citrine/lolo/validation/MeritTest.scala
index 51dcd4e5..bdf01546 100644
--- a/src/test/scala/io/citrine/lolo/validation/MeritTest.scala
+++ b/src/test/scala/io/citrine/lolo/validation/MeritTest.scala
@@ -1,7 +1,8 @@
 package io.citrine.lolo.validation
 
 import io.citrine.random.Random
-import io.citrine.lolo.{PredictionResult, SeedRandomMixIn}
+import io.citrine.lolo.SeedRandomMixIn
+import io.citrine.lolo.api.PredictionResult
 import org.apache.commons.math3.distribution.MultivariateNormalDistribution
 import org.apache.commons.math3.random.MersenneTwister
 import org.junit.Test
@@ -54,9 +55,9 @@ class MeritTest extends SeedRandomMixIn {
         (y + error, uncertainty, y)
       }
       val predictionResult = new PredictionResult[Double] {
-        override def getExpected(): Seq[Double] = pua.map(_._1)
+        override def expected: Seq[Double] = pua.map(_._1)
 
-        override def getUncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = Some(pua.map(_._2))
+        override def uncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = Some(pua.map(_._2))
       }
       (predictionResult, pua.map(_._3))
     }
diff --git a/src/test/scala/io/citrine/lolo/validation/StatisticalValidationTest.scala b/src/test/scala/io/citrine/lolo/validation/StatisticalValidationTest.scala
index 6e34575b..c2bde39f 100644
--- a/src/test/scala/io/citrine/lolo/validation/StatisticalValidationTest.scala
+++ b/src/test/scala/io/citrine/lolo/validation/StatisticalValidationTest.scala
@@ -1,7 +1,7 @@
 package io.citrine.lolo.validation
 
-import io.citrine.lolo.{SeedRandomMixIn, TestUtils}
-import io.citrine.lolo.learners.RandomForest
+import io.citrine.lolo.{DataGenerator, SeedRandomMixIn}
+import io.citrine.lolo.learners.RandomForestRegressor
 import io.citrine.lolo.stats.functions.Friedman
 import org.junit.Test
 
@@ -9,9 +9,9 @@ class StatisticalValidationTest extends SeedRandomMixIn {
 
   @Test
   def testCompareToKFolds(): Unit = {
-    val learner = RandomForest()
-    val dataSet = TestUtils.generateTrainingData(128, 8, Friedman.friedmanSilverman, rng = rng)
-    val dataGenerator = TestUtils.iterateTrainingData(8, Friedman.friedmanSilverman, rng = rng)
+    val learner = RandomForestRegressor()
+    val dataSet = DataGenerator.generate(128, 8, Friedman.friedmanSilverman, rng = rng).data
+    val dataGenerator = DataGenerator.iterate(8, Friedman.friedmanSilverman, rng = rng)
 
     val metrics = Map("rmse" -> RootMeanSquareError)
     val (rmseFromCV, uncertaintyFromCV) =