From 3fca5d4bf08f1d378d88c585dfa854f93e3fde30 Mon Sep 17 00:00:00 2001
From: adpauls <adpauls@semanticmachines.com>
Date: Thu, 26 May 2016 14:18:58 -0700
Subject: [PATCH 1/7] Cleanup.

---
 src/main/scala/epic/logo/Trainer.scala | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/main/scala/epic/logo/Trainer.scala b/src/main/scala/epic/logo/Trainer.scala
index db23f627..8c6682e0 100644
--- a/src/main/scala/epic/logo/Trainer.scala
+++ b/src/main/scala/epic/logo/Trainer.scala
@@ -252,12 +252,14 @@ case class Trainer[T, W](convergenceChecker: ConvergenceChecker[W],
         }
       }
       loopWhile(numOuterOptimizationLoops) {
-        data.exists { instance =>
+        // Note: we don't use exists because we want to run on every example, not stop on the first
+        // one with a change
+        !data.forall { instance =>
           val numUpdatesExecuted = loopWhile(numInnerOptimizationLoops) {
             updater.update(instance, w, n, iteration)
           }
-          // > 1 instead of > 0 because update is always called once, but it might do nothing.
-          numUpdatesExecuted > 1
+          // 1 instead of 0 because update is always called once, but it might do nothing.
+          numUpdatesExecuted <= 1
         }
       }
       iterationCallback.endIteration(iteration, w)

From 350b2325a98df13c4314b311f9d73e0d8c48c888 Mon Sep 17 00:00:00 2001
From: adpauls <adpauls@semanticmachines.com>
Date: Thu, 26 May 2016 14:46:41 -0700
Subject: [PATCH 2/7] CLeanup.

---
 src/main/scala/epic/logo/Trainer.scala | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/main/scala/epic/logo/Trainer.scala b/src/main/scala/epic/logo/Trainer.scala
index 8c6682e0..bdfcf4ce 100644
--- a/src/main/scala/epic/logo/Trainer.scala
+++ b/src/main/scala/epic/logo/Trainer.scala
@@ -157,12 +157,10 @@ object Trainer {
       val trainer = newL1MaxMarginTrainer(new MulticlassOneSlackOracleInferencer(oracler),
         new MulticlassOneSlackLossAugmentedArgmaxInferencer(argmaxer, initialConstraint),
         oneSlackIterCallBack, C, maxNumIters, opts)
-      trainer.train(Seq(new Instance(data)),
-        new Weights(space.zeroLike(initialConstraint)))
+      trainer.train(Seq(data), new Weights(space.zeroLike(initialConstraint)))
     } else {
       val trainer = newL1MaxMarginTrainer(oracler, argmaxer, iterationCallback, C, maxNumIters, opts)
-      trainer.train(data.map(new Instance[LabeledDatum[L, F], W](_)),
-        new Weights(space.zeroLike(initialConstraint)))
+      trainer.train(data, new Weights(space.zeroLike(initialConstraint)))
     }
     new MulticlassClassifier[L, F, W](weights, argmaxer)
   }
@@ -182,7 +180,8 @@ case class Trainer[T, W](convergenceChecker: ConvergenceChecker[W],
   final val numInnerOptimizationLoops = if (online) 1 else opts.numInnerOptimizationLoops
   final val numOuterOptimizationLoops = if (online) 0 else opts.numOuterOptimizationLoops
 
-  def train(data : Seq[Instance[T, W]], initWeights : Weights[W]) = {
+  def train(rawData : Seq[T], initWeights : Weights[W]) = {
+    val data = rawData.map(datum => Instance[T, W](datum))
     val n = data.length
     val w = if (online) initWeights else {
       if (initWeights.norm != 0.0) {

From 8790455eefc795c88d0f39fe461f943e3f0c217f Mon Sep 17 00:00:00 2001
From: adpauls <adpauls@semanticmachines.com>
Date: Thu, 26 May 2016 15:39:02 -0700
Subject: [PATCH 3/7] More cleanup

---
 src/main/scala/epic/logo/Trainer.scala | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/main/scala/epic/logo/Trainer.scala b/src/main/scala/epic/logo/Trainer.scala
index bdfcf4ce..97cda4b5 100644
--- a/src/main/scala/epic/logo/Trainer.scala
+++ b/src/main/scala/epic/logo/Trainer.scala
@@ -253,13 +253,14 @@ case class Trainer[T, W](convergenceChecker: ConvergenceChecker[W],
       loopWhile(numOuterOptimizationLoops) {
         // Note: we don't use exists because we want to run on every example, not stop on the first
         // one with a change
-        !data.forall { instance =>
+        val instancesChanged = data.count { instance =>
           val numUpdatesExecuted = loopWhile(numInnerOptimizationLoops) {
             updater.update(instance, w, n, iteration)
           }
           // 1 instead of 0 because update is always called once, but it might do nothing.
           numUpdatesExecuted <= 1
         }
+        instancesChanged > 0
       }
       iterationCallback.endIteration(iteration, w)
       if (average) {

From 6ecc5545088a0fba78cfb9a5eb121a483da5826f Mon Sep 17 00:00:00 2001
From: adpauls <adpauls@semanticmachines.com>
Date: Thu, 26 May 2016 15:50:15 -0700
Subject: [PATCH 4/7] Fix

---
 .../scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala b/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala
index 379ddd69..ab6cb2b4 100644
--- a/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala
+++ b/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala
@@ -13,7 +13,7 @@ case class ObjectiveFunctionConvergenceChecker[W](
   private def objectiveConverged(weights: Weights[W], data: Seq[Instance[_, W]], iter: Int) = {
     val (primal, dual) = objective.calculatePrimalAndDual(weights, data)
     callback.objectiveValCheck(primal, dual, iter, weights)
-    NumUtils.approxEquals(primal, dual, 1e-10)
+    NumUtils.approxEquals(primal, dual, tol)
   }
 
 }

From eb01b2fce81b0be6da92d1cb939eff646145697a Mon Sep 17 00:00:00 2001
From: adpauls <adpauls@semanticmachines.com>
Date: Thu, 26 May 2016 18:19:13 -0700
Subject: [PATCH 5/7] Give decoders the ability to track state across calls.

---
 .../scala/epic/logo/ArgmaxInferencer.scala    |   8 +-
 .../epic/logo/CompoundIterationCallback.scala |  20 ++--
 src/main/scala/epic/logo/Decoder.scala        |   7 +-
 .../epic/logo/ExpectationInferencer.scala     |   5 +-
 src/main/scala/epic/logo/Inferencer.scala     |   7 ++
 .../scala/epic/logo/IterationCallback.scala   |   8 +-
 .../epic/logo/LogLikelihoodDecoder.scala      |  18 +--
 .../logo/LossAugmentedArgmaxInferencer.scala  |   6 +-
 .../LossAugmentedExpectationInferencer.scala  |   6 +-
 .../logo/LossAugmentedMaxMarginDecoder.scala  |  18 ++-
 .../scala/epic/logo/MaxMarginDecoder.scala    |  18 ++-
 .../epic/logo/MaxMarginRankingDecoder.scala   |  19 +++-
 src/main/scala/epic/logo/MinibatchInput.scala |   6 +-
 ...ticlassLossAugmentedArgmaxInferencer.scala |  12 +-
 ...neSlackLossAugmentedArgmaxInferencer.scala |  16 ++-
 .../MulticlassOneSlackOracleInferencer.scala  |  12 +-
 .../logo/MulticlassOracleInferencer.scala     |  10 +-
 .../ObjectiveFunctionConvergenceChecker.scala |   2 +-
 .../scala/epic/logo/OracleInferencer.scala    |   6 +-
 src/main/scala/epic/logo/Trainer.scala        | 107 ++++++++++--------
 20 files changed, 194 insertions(+), 117 deletions(-)
 create mode 100644 src/main/scala/epic/logo/Inferencer.scala

diff --git a/src/main/scala/epic/logo/ArgmaxInferencer.scala b/src/main/scala/epic/logo/ArgmaxInferencer.scala
index c105f57a..61804767 100644
--- a/src/main/scala/epic/logo/ArgmaxInferencer.scala
+++ b/src/main/scala/epic/logo/ArgmaxInferencer.scala
@@ -1,7 +1,9 @@
 package epic.logo
 
-trait ArgmaxInferencer[T, Y, W] {
-
-  def argmax(weights : Weights[W], instance : T) : (Y, W, Double)
+trait ArgmaxInferencer[T, W, S] extends Inferencer[S] {
+  type Y
 
+  def argmax(weights : Weights[W], instance : T) : (Y, W, Double, S)
+  def initialState: S
+  def reduceStates(s1: S, s2: S): S
 }
diff --git a/src/main/scala/epic/logo/CompoundIterationCallback.scala b/src/main/scala/epic/logo/CompoundIterationCallback.scala
index 1f3b0cb5..aaa746f5 100644
--- a/src/main/scala/epic/logo/CompoundIterationCallback.scala
+++ b/src/main/scala/epic/logo/CompoundIterationCallback.scala
@@ -1,27 +1,31 @@
 package epic.logo
 
-class CompoundIterationCallback[T, W](val callbacks : Iterable[IterationCallback[T, W]]) extends IterationCallback[T, W] {
+class CompoundIterationCallback[T, W, S1, S2](val callbacks: Iterable[IterationCallback[T, W, S1, S2]])
+    extends IterationCallback[T, W, S1, S2] {
 
-  override def startIteration(iter : Int, weights : Weights[W]) : Unit = {
+  override def startIteration(iter: Int, weights: Weights[W]): Unit = {
     callbacks.foreach(_.startIteration(iter, weights))
   }
 
-  override def startMinibatch(iter : Int, weights : Weights[W], miniBatch : Array[MinibatchInput[T, W]]) : Unit = {
+  override def startMinibatch(iter: Int, weights: Weights[W],
+                              miniBatch: Array[MinibatchInput[T, W]]): Unit = {
     callbacks.foreach(_.startMinibatch(iter, weights, miniBatch))
   }
-  override def endMinibatch(iter : Int, weights : Weights[W], miniBatch : Array[MinibatchOutput[T, W]]) : Unit = {
+  override def endMinibatch(iter: Int, weights: Weights[W],
+                            miniBatch: Array[MinibatchOutput[T, W, S1, S2]]): Unit = {
     callbacks.foreach(_.endMinibatch(iter, weights, miniBatch))
   }
 
-  override def endIteration(iter : Int, weights : Weights[W]) : Unit = {
-    callbacks.foreach(_.endIteration(iter, weights))
+  override def endIteration(iter: Int, weights: Weights[W], state1: S1, state2: S2): Unit = {
+    callbacks.foreach(_.endIteration(iter, weights, state1, state2))
   }
 
-  override def objectiveValCheck(primal : Double, dual : Double, iter : Int, w : Weights[W]) : Unit = {
+  override def objectiveValCheck(primal: Double, dual: Double, iter: Int, w: Weights[W]): Unit = {
     callbacks.foreach(_.objectiveValCheck(primal, dual, iter, w))
   }
 
-  override def converged(weights : Weights[W], data : Seq[Instance[T, W]], iter : Int, numNewConstraints : Int) : Boolean = {
+  override def converged(weights: Weights[W], data: Seq[Instance[T, W]],
+                         iter: Int, numNewConstraints: Int): Boolean = {
     callbacks.forall(_.converged(weights, data, iter, numNewConstraints))
   }
 
diff --git a/src/main/scala/epic/logo/Decoder.scala b/src/main/scala/epic/logo/Decoder.scala
index 00e6fbfa..65eb3e02 100644
--- a/src/main/scala/epic/logo/Decoder.scala
+++ b/src/main/scala/epic/logo/Decoder.scala
@@ -1,7 +1,8 @@
 package epic.logo
 
-trait Decoder[T, W] {
-
-  def decode(weights : Weights[W], instance : T) : (W, Double)
+trait Decoder[T, W, OracleS, MaxerS] {
 
+  def decode(weights : Weights[W], instance : T) : (W, Double, OracleS, MaxerS)
+  def initialState: (OracleS, MaxerS)
+  def reduceStates(states1: (OracleS, MaxerS), states2: (OracleS, MaxerS)): (OracleS, MaxerS)
 }
diff --git a/src/main/scala/epic/logo/ExpectationInferencer.scala b/src/main/scala/epic/logo/ExpectationInferencer.scala
index 2a53adbf..9d54f4ed 100644
--- a/src/main/scala/epic/logo/ExpectationInferencer.scala
+++ b/src/main/scala/epic/logo/ExpectationInferencer.scala
@@ -1,7 +1,6 @@
 package epic.logo
 
-trait ExpectationInferencer[T, W] {
-
-  def expectations(weights : Weights[W], instance : T) : (W, Double)
+trait ExpectationInferencer[T, W, S] extends Inferencer[S] {
 
+  def expectations(weights: Weights[W], instance: T): (W, Double, S)
 }
diff --git a/src/main/scala/epic/logo/Inferencer.scala b/src/main/scala/epic/logo/Inferencer.scala
new file mode 100644
index 00000000..5ad680be
--- /dev/null
+++ b/src/main/scala/epic/logo/Inferencer.scala
@@ -0,0 +1,7 @@
+package epic.logo
+
+trait Inferencer[S] {
+  def initialState: S
+  def reduceStates(state1: S, state2: S): S
+
+}
diff --git a/src/main/scala/epic/logo/IterationCallback.scala b/src/main/scala/epic/logo/IterationCallback.scala
index d493c711..1185a25e 100644
--- a/src/main/scala/epic/logo/IterationCallback.scala
+++ b/src/main/scala/epic/logo/IterationCallback.scala
@@ -1,12 +1,12 @@
 package epic.logo
 
-trait IterationCallback[T, W] {
+trait IterationCallback[T, W, OracleS, MaxerS] {
 
   def startIteration(iter : Int, weights : Weights[W]) : Unit = {}
 
   def startMinibatch(iter : Int, weights : Weights[W], miniBatch : Array[MinibatchInput[T, W]]) : Unit = {}
-  def endMinibatch(iter : Int, weights : Weights[W], miniBatch : Array[MinibatchOutput[T, W]]) : Unit = {}
-  def endIteration(iter : Int, weights : Weights[W]) : Unit = {}
+  def endMinibatch(iter : Int, weights : Weights[W], miniBatch : Array[MinibatchOutput[T, W, OracleS, MaxerS]]) : Unit = {}
+  def endIteration(iter : Int, weights : Weights[W], oracleState: OracleS, maxerState: MaxerS) : Unit = {}
 
   def objectiveValCheck(primal : Double, dual : Double, iter : Int, weights : Weights[W]) : Unit = {}
 
@@ -14,4 +14,4 @@ trait IterationCallback[T, W] {
 
 }
 
-case class NullIterationCallback[T, W]() extends IterationCallback[T, W]
+case class NullIterationCallback[T, W, S1, S2]() extends IterationCallback[T, W, S1, S2]
diff --git a/src/main/scala/epic/logo/LogLikelihoodDecoder.scala b/src/main/scala/epic/logo/LogLikelihoodDecoder.scala
index f157cfcf..7c6f603a 100644
--- a/src/main/scala/epic/logo/LogLikelihoodDecoder.scala
+++ b/src/main/scala/epic/logo/LogLikelihoodDecoder.scala
@@ -2,15 +2,19 @@ package epic.logo
 
 import breeze.math.MutableInnerProductModule
 
-class LogLikelihoodDecoder[T, W](
-    val inferencer: OracleInferencer[T, _, W], val summer: ExpectationInferencer[T, W])(
-        implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W] {
+case class LogLikelihoodDecoder[T, W, OracleS, ExpectationS](
+    inferencer: OracleInferencer[T, W, OracleS], summer: ExpectationInferencer[T, W, ExpectationS])(
+        implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W, OracleS, ExpectationS] {
   import space._
 
-  def decode(weights : Weights[W], instance : T) : (W, Double) = {
-    val (y_*, f_*, l_*) = inferencer.oracle(weights, instance)
-    val (f, l) = summer.expectations(weights, instance)
-    ((f_* - f), l - l_*)
+  def decode(weights : Weights[W], instance : T) : (W, Double, OracleS, ExpectationS) = {
+    val (y_*, f_*, l_*, state_*) = inferencer.oracle(weights, instance)
+    val (f, l, state) = summer.expectations(weights, instance)
+    ((f_* - f), l - l_*, state_*, state)
   }
 
+  def initialState = (inferencer.initialState, summer.initialState)
+  def reduceStates(states1: (OracleS, ExpectationS), states2: (OracleS, ExpectationS)): (OracleS, ExpectationS) =
+    (inferencer.reduceStates(states1._1, states2._1),
+      summer.reduceStates(states1._2, states2._2))
 }
diff --git a/src/main/scala/epic/logo/LossAugmentedArgmaxInferencer.scala b/src/main/scala/epic/logo/LossAugmentedArgmaxInferencer.scala
index 54b3096a..309da55b 100644
--- a/src/main/scala/epic/logo/LossAugmentedArgmaxInferencer.scala
+++ b/src/main/scala/epic/logo/LossAugmentedArgmaxInferencer.scala
@@ -1,9 +1,9 @@
 package epic.logo
 
-trait LossAugmentedArgmaxInferencer[T, Y, W] extends ArgmaxInferencer[T, Y, W] {
+trait LossAugmentedArgmaxInferencer[T, W, S] extends ArgmaxInferencer[T, W, S] {
 
-  def argmax(weights : Weights[W], instance : T) : (Y, W, Double) = lossAugmentedArgmax(weights, instance, 1.0, 0.0)
+  def argmax(weights : Weights[W], instance : T) : (Y, W, Double, S) = lossAugmentedArgmax(weights, instance, 1.0, 0.0)
 
-  def lossAugmentedArgmax(weights : Weights[W], instance : T, weightsWeight : Double, lossWeight : Double) : (Y, W, Double)
+  def lossAugmentedArgmax(weights : Weights[W], instance : T, weightsWeight : Double, lossWeight : Double) : (Y, W, Double, S)
 
 }
diff --git a/src/main/scala/epic/logo/LossAugmentedExpectationInferencer.scala b/src/main/scala/epic/logo/LossAugmentedExpectationInferencer.scala
index 86d7db13..3b04fa0e 100644
--- a/src/main/scala/epic/logo/LossAugmentedExpectationInferencer.scala
+++ b/src/main/scala/epic/logo/LossAugmentedExpectationInferencer.scala
@@ -1,9 +1,9 @@
 package epic.logo
 
-trait LossAugmentedExpectationInferencer[T, W] extends ExpectationInferencer[T, W] {
+trait LossAugmentedExpectationInferencer[T, W,  S] extends ExpectationInferencer[T, W, S] {
 
-  def expectations(weights : Weights[W], instance : T) : (W, Double) = lossAugmentedExpectations(weights, instance, 1.0, 0.0)
+  def expectations(weights : Weights[W], instance : T) : (W, Double, S) = lossAugmentedExpectations(weights, instance, 1.0, 0.0)
 
-  def lossAugmentedExpectations(weights : Weights[W], instance : T, weightsWeight : Double, lossWeight : Double) : (W, Double)
+  def lossAugmentedExpectations(weights : Weights[W], instance : T, weightsWeight : Double, lossWeight : Double) : (W, Double, S)
 
 }
diff --git a/src/main/scala/epic/logo/LossAugmentedMaxMarginDecoder.scala b/src/main/scala/epic/logo/LossAugmentedMaxMarginDecoder.scala
index 16f23438..810002a8 100644
--- a/src/main/scala/epic/logo/LossAugmentedMaxMarginDecoder.scala
+++ b/src/main/scala/epic/logo/LossAugmentedMaxMarginDecoder.scala
@@ -2,13 +2,21 @@ package epic.logo
 
 import breeze.math.MutableInnerProductModule
 
-class LossAugmentedMaxMarginDecoder[T, W](val oracleInferencer: OracleInferencer[T, _, W], val argmaxer: LossAugmentedArgmaxInferencer[T, _, W])(implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W] {
+case class LossAugmentedMaxMarginDecoder[T, W, OracleS, ArgmaxerS](
+    oracleInferencer: OracleInferencer[T, W, OracleS],
+    argmaxer: LossAugmentedArgmaxInferencer[T, W, ArgmaxerS])(implicit space: MutableInnerProductModule[W, Double])
+    extends Decoder[T, W, OracleS, ArgmaxerS] {
   import space._
 
-  def decode(weights: Weights[W], instance: T): (W, Double) = {
-    val (y_*, f_*, l_*) = oracleInferencer.oracle(weights, instance)
-    val (y, f, l) = argmaxer.lossAugmentedArgmax(weights, instance, 1.0, 1.0)
-    (f_* - f, l - l_*)
+  def decode(weights: Weights[W], instance: T): (W, Double, OracleS, ArgmaxerS) = {
+    val (y_*, f_*, l_*, state_*) = oracleInferencer.oracle(weights, instance)
+    val (y, f, l, state) = argmaxer.lossAugmentedArgmax(weights, instance, 1.0, 1.0)
+    (f_* - f, l - l_*, state_*, state)
   }
 
+  def initialState: (OracleS, ArgmaxerS) = (oracleInferencer.initialState, argmaxer.initialState)
+  def reduceStates(states1: (OracleS, ArgmaxerS), states2: (OracleS, ArgmaxerS)): (OracleS, ArgmaxerS) =
+    (oracleInferencer.reduceStates(states1._1, states2._1),
+      argmaxer.reduceStates(states1._2, states2._2))
+
 }
diff --git a/src/main/scala/epic/logo/MaxMarginDecoder.scala b/src/main/scala/epic/logo/MaxMarginDecoder.scala
index 80956e54..504acd51 100644
--- a/src/main/scala/epic/logo/MaxMarginDecoder.scala
+++ b/src/main/scala/epic/logo/MaxMarginDecoder.scala
@@ -2,13 +2,21 @@ package epic.logo
 
 import breeze.math.MutableInnerProductModule
 
-class MaxMarginDecoder[T, W](val oracleInferencer : OracleInferencer[T, _, W], val argmaxer : ArgmaxInferencer[T, _, W])(implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W] {
+case class MaxMarginDecoder[T, W, OracleS, MaxerS](
+  oracleInferencer: OracleInferencer[T, W, OracleS],
+  argmaxer: ArgmaxInferencer[T, W, MaxerS])(implicit space: MutableInnerProductModule[W, Double])
+    extends Decoder[T, W, OracleS, MaxerS] {
   import space._
 
-  def decode(weights : Weights[W], instance : T) : (W, Double) = {
-    val (y_*, f_*, l_*) = oracleInferencer.oracle(weights, instance)
-    val (y, f, l) = argmaxer.argmax(weights, instance)
-    (f_* - f, l - l_*)
+  def decode(weights : Weights[W], instance : T) : (W, Double, OracleS, MaxerS) = {
+    val (y_*, f_*, l_*, state_*) = oracleInferencer.oracle(weights, instance)
+    val (y, f, l, state) = argmaxer.argmax(weights, instance)
+    (f_* - f, l - l_*, state_*, state)
   }
 
+  def initialState: (OracleS, MaxerS) = (oracleInferencer.initialState, argmaxer.initialState)
+  def reduceStates(states1: (OracleS, MaxerS), states2: (OracleS, MaxerS)): (OracleS, MaxerS) =
+    (oracleInferencer.reduceStates(states1._1, states2._1),
+      argmaxer.reduceStates(states1._2, states2._2))
+
 }
diff --git a/src/main/scala/epic/logo/MaxMarginRankingDecoder.scala b/src/main/scala/epic/logo/MaxMarginRankingDecoder.scala
index fee00c01..31304853 100644
--- a/src/main/scala/epic/logo/MaxMarginRankingDecoder.scala
+++ b/src/main/scala/epic/logo/MaxMarginRankingDecoder.scala
@@ -2,12 +2,19 @@ package epic.logo
 
 import breeze.math.MutableInnerProductModule
 
-class MaxMarginRankingDecoder[T, W](val inferencer : LossAugmentedArgmaxInferencer[T, _, W], val gamma : Double = 0.0)(implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W] {
-import space._
-  def decode(weights: Weights[W], instance: T): (W, Double) = {
-    val (y_min, df_min, l_min) = inferencer.lossAugmentedArgmax(weights, instance, -1.0, -(1.0 + gamma))
-    val (y_max, df_max, l_max) = inferencer.lossAugmentedArgmax(weights, instance, 1.0, 1.0)
-    (df_min - df_max, l_max - (1.0 + gamma) * l_min)
+case class MaxMarginRankingDecoder[T, W, S](
+    inferencer: LossAugmentedArgmaxInferencer[T, W, S],
+    gamma: Double = 0.0)(implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W, S, S] {
+  import space._
+  def decode(weights: Weights[W], instance: T): (W, Double, S, S) = {
+    val (y_min, df_min, l_min, s_min) = inferencer.lossAugmentedArgmax(weights, instance, -1.0, -(1.0 + gamma))
+    val (y_max, df_max, l_max, s_max) = inferencer.lossAugmentedArgmax(weights, instance, 1.0, 1.0)
+    (df_min - df_max, l_max - (1.0 + gamma) * l_min, s_min, s_max)
   }
 
+  def initialState: (S, S) = (inferencer.initialState, inferencer.initialState)
+  def reduceStates(states1: (S, S), states2: (S, S)): (S, S) =
+    (inferencer.reduceStates(states1._1, states2._1),
+      inferencer.reduceStates(states1._2, states2._2))
+
 }
diff --git a/src/main/scala/epic/logo/MinibatchInput.scala b/src/main/scala/epic/logo/MinibatchInput.scala
index 0035f066..d6456012 100644
--- a/src/main/scala/epic/logo/MinibatchInput.scala
+++ b/src/main/scala/epic/logo/MinibatchInput.scala
@@ -1,5 +1,7 @@
 package epic.logo
 
-case class MinibatchInput[T, W](val instance : Instance[T, W], val instanceNum : Int)
+case class MinibatchInput[T, W](instance : Instance[T, W], instanceNum : Int)
 
-case class MinibatchOutput[T, W](val instance : Instance[T, W], val instanceNum : Int, val df : W, val loss : Double)
+case class MinibatchOutput[T, W, OracleS, MaxerS](
+  instance: Instance[T, W], instanceNum: Int, df: W, loss: Double,
+  oracleState: OracleS, maxerState: MaxerS)
diff --git a/src/main/scala/epic/logo/MulticlassLossAugmentedArgmaxInferencer.scala b/src/main/scala/epic/logo/MulticlassLossAugmentedArgmaxInferencer.scala
index 84e5cac1..f7b08bdc 100644
--- a/src/main/scala/epic/logo/MulticlassLossAugmentedArgmaxInferencer.scala
+++ b/src/main/scala/epic/logo/MulticlassLossAugmentedArgmaxInferencer.scala
@@ -3,16 +3,18 @@ import breeze.util.Index
 import breeze.math.MutableInnerProductModule
 
 class MulticlassLossAugmentedArgmaxInferencer[L, F, W](validLabels: IndexedSeq[L], labelConjoiner: (L, F) => W)
-(implicit space: MutableInnerProductModule[W, Double])
-    extends LossAugmentedArgmaxInferencer[LabeledDatum[L, F], L, W] {
+    extends LossAugmentedArgmaxInferencer[LabeledDatum[L, F], W, Unit] {
+  type Y = L
 
   def lossAugmentedArgmax(weights: Weights[W], instance: LabeledDatum[L, F],
-                          weightsWeight: Double, lossWeight: Double): (L, W, Double) = {
+                          weightsWeight: Double, lossWeight: Double): (L, W, Double, Unit) = {
     validLabels.map(label => {
       val loss = if (instance.label == null || label.equals(instance.label)) 0.0 else 1.0
       val labeledFeatureVector = labelConjoiner(label, instance.features)
-      (label,labeledFeatureVector, loss)
-    }).maxBy { case (label, features, loss) => weightsWeight * (weights * features) + lossWeight * loss }
+      (label,labeledFeatureVector, loss, ())
+    }).maxBy { case (label, features, loss, _) => weightsWeight * (weights * features) + lossWeight * loss }
   }
 
+  def initialState: Unit = ()
+  def reduceStates(a: Unit, b: Unit): Unit = ()
 }
diff --git a/src/main/scala/epic/logo/MulticlassOneSlackLossAugmentedArgmaxInferencer.scala b/src/main/scala/epic/logo/MulticlassOneSlackLossAugmentedArgmaxInferencer.scala
index 0fd39cc6..68871084 100644
--- a/src/main/scala/epic/logo/MulticlassOneSlackLossAugmentedArgmaxInferencer.scala
+++ b/src/main/scala/epic/logo/MulticlassOneSlackLossAugmentedArgmaxInferencer.scala
@@ -3,16 +3,24 @@ package epic.logo
 import breeze.math.MutableInnerProductModule
 
 class MulticlassOneSlackLossAugmentedArgmaxInferencer[L, F, W](
-    inferencer: MulticlassLossAugmentedArgmaxInferencer[L, F, W],
-    emptyFeatureVector: W)(implicit space: MutableInnerProductModule[W, Double]) extends LossAugmentedArgmaxInferencer[Seq[LabeledDatum[L, F]], Seq[L], W] {
+  inferencer: MulticlassLossAugmentedArgmaxInferencer[L, F, W],
+  emptyFeatureVector: W)(implicit space: MutableInnerProductModule[W, Double])
+    extends LossAugmentedArgmaxInferencer[Seq[LabeledDatum[L, F]], W, Unit] {
+  override type Y = Seq[L]
   import space._
 
   def lossAugmentedArgmax(weights: Weights[W], instance: Seq[LabeledDatum[L, F]], weightsWeight: Double,
-                          lossWeight: Double): (Seq[L], W, Double) = {
-    instance.map(i => inferencer.lossAugmentedArgmax(weights, i, weightsWeight, lossWeight))
+                          lossWeight: Double): (Seq[L], W, Double, Unit) = {
+   val (labels, fvs, losses) =
+     instance.map(i => inferencer.lossAugmentedArgmax(weights, i, weightsWeight, lossWeight))
       .foldLeft((Seq.empty[L], emptyFeatureVector, 0.0)) { (acc, curr) =>
         (acc._1 :+ curr._1, acc._2 + curr._2, acc._3 + curr._3)
       }
+   (labels, fvs, losses, ())
   }
 
+  def initialState: Unit = ()
+
+  def reduceStates(s1: Unit, s2: Unit): Unit = ()
+
 }
diff --git a/src/main/scala/epic/logo/MulticlassOneSlackOracleInferencer.scala b/src/main/scala/epic/logo/MulticlassOneSlackOracleInferencer.scala
index 3a20e051..3f1b76fe 100644
--- a/src/main/scala/epic/logo/MulticlassOneSlackOracleInferencer.scala
+++ b/src/main/scala/epic/logo/MulticlassOneSlackOracleInferencer.scala
@@ -4,15 +4,21 @@ import breeze.math.MutableInnerProductModule
 
 class MulticlassOneSlackOracleInferencer[L, F, W](
   val inferencer: MulticlassOracleInferencer[L, F, W])(implicit space: MutableInnerProductModule[W, Double])
-    extends OracleInferencer[Seq[LabeledDatum[L, F]], Seq[L], W] {
+    extends OracleInferencer[Seq[LabeledDatum[L, F]], W, Unit] {
   import space._
+  type Y = Seq[L]
 
-  def oracle(weights: Weights[W], instance: Seq[LabeledDatum[L, F]]): (Seq[L], W, Double) = {
+  def oracle(weights: Weights[W], instance: Seq[LabeledDatum[L, F]]): (Seq[L], W, Double, Unit) = {
     val zeroVector = space.zeroLike(inferencer.labelConjoiner(instance.head.label, instance.head.features))
-    instance.map(i => inferencer.oracle(weights, i))
+    val (labels, fvs, losses) = instance.map(i => inferencer.oracle(weights, i))
       .foldLeft((Seq.empty[L], zeroVector, 0.0)) { (acc, curr) =>
         (acc._1 :+ curr._1, acc._2 + curr._2, acc._3 + curr._3)
       }
+    (labels, fvs, losses, ())
   }
 
+  def initialState: Unit = ()
+
+  def reduceStates(state1: Unit, state2: Unit): Unit = ()
+
 }
diff --git a/src/main/scala/epic/logo/MulticlassOracleInferencer.scala b/src/main/scala/epic/logo/MulticlassOracleInferencer.scala
index 25345e79..05404c18 100644
--- a/src/main/scala/epic/logo/MulticlassOracleInferencer.scala
+++ b/src/main/scala/epic/logo/MulticlassOracleInferencer.scala
@@ -5,11 +5,15 @@ import breeze.math.MutableInnerProductModule
 
 case class MulticlassOracleInferencer[L, F, W](
   validLabels: IndexedSeq[L], labelConjoiner: (L, F) => W)(implicit space: MutableInnerProductModule[W, Double])
-    extends OracleInferencer[LabeledDatum[L, F], L, W] {
+    extends OracleInferencer[LabeledDatum[L, F], W, Unit] {
+  override type Y = L
 
-  def oracle(weights : Weights[W], instance : LabeledDatum[L, F]) : (L, W, Double) = {
-    (instance.label, labelConjoiner(instance.label, instance.features), 0.0)
+  def oracle(weights : Weights[W], instance : LabeledDatum[L, F]) : (L, W, Double, Unit) = {
+    (instance.label, labelConjoiner(instance.label, instance.features), 0.0, ())
 
   }
 
+  def initialState = ()
+  def reduceStates(a: Unit, b: Unit) = ()
+
 }
diff --git a/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala b/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala
index ab6cb2b4..cddc2f15 100644
--- a/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala
+++ b/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala
@@ -4,7 +4,7 @@ import scala.collection.Seq
 
 case class ObjectiveFunctionConvergenceChecker[W](
     objective: ObjectiveFunction[W], maxNumIters: Int,
-    callback: IterationCallback[_, W], tol: Double) extends ConvergenceChecker[W] {
+    callback: IterationCallback[_, W, _, _], tol: Double) extends ConvergenceChecker[W] {
 
   def converged(weights: Weights[W], data: Seq[Instance[_, W]], iter: Int, numNewConstraints: Int): Boolean = {
     iter >= maxNumIters || (objectiveConverged(weights, data, iter) && numNewConstraints == 0)
diff --git a/src/main/scala/epic/logo/OracleInferencer.scala b/src/main/scala/epic/logo/OracleInferencer.scala
index 532f7236..f3710cd4 100644
--- a/src/main/scala/epic/logo/OracleInferencer.scala
+++ b/src/main/scala/epic/logo/OracleInferencer.scala
@@ -1,7 +1,7 @@
 package epic.logo
 
-trait OracleInferencer[T, Y, W] {
-
-  def oracle(weights : Weights[W], instance : T) : (Y, W, Double)
+trait OracleInferencer[T, W, S] extends Inferencer[S] {
+  type Y
 
+  def oracle(weights : Weights[W], instance : T) : (Y, W, Double, S)
 }
diff --git a/src/main/scala/epic/logo/Trainer.scala b/src/main/scala/epic/logo/Trainer.scala
index 97cda4b5..6b93a500 100644
--- a/src/main/scala/epic/logo/Trainer.scala
+++ b/src/main/scala/epic/logo/Trainer.scala
@@ -12,9 +12,9 @@ import breeze.math.MutableInnerProductModule
 import scala.reflect.ClassTag
 
 object Trainer {
-  def newL1MaxMarginTrainer[T, Y, W](
-    oracleInferencer: OracleInferencer[T, Y, W], argmaxer: LossAugmentedArgmaxInferencer[T, Y, W],
-    iterationCallback: IterationCallback[T, W] = NullIterationCallback[T, W](),
+  def newL1MaxMarginTrainer[T, W, OracleS, MaxerS](
+    oracleInferencer: OracleInferencer[T, W, OracleS], argmaxer: LossAugmentedArgmaxInferencer[T, W, MaxerS],
+    iterationCallback: IterationCallback[T, W, OracleS, MaxerS] = NullIterationCallback(),
     C: Double = 1.0,
     maxNumIters: Int = 100, opts: LogoOpts = new LogoOpts(),
     addInitialConstraint: Option[W] = None)(implicit space: MutableInnerProductModule[W, Double]) =
@@ -27,9 +27,9 @@ object Trainer {
         initialConstraintAndAlpha = addInitialConstraint.map(c => ((c, 0.0), C)))
     }
 
-  def newL1MIRATrainer[T, Y, W](
-    oracleInferencer: OracleInferencer[T, Y, W], argmaxer: LossAugmentedArgmaxInferencer[T, Y, W],
-    iterationCallback: IterationCallback[T, W], C: Double = 1.0,
+  def newL1MIRATrainer[T, Y, W, OracleS, MaxerS](
+    oracleInferencer: OracleInferencer[T, W, OracleS], argmaxer: LossAugmentedArgmaxInferencer[T, W, MaxerS],
+    iterationCallback: IterationCallback[T, W, OracleS, MaxerS], C: Double = 1.0,
     maxNumIters: Int = 100, opts: LogoOpts = new LogoOpts(), average: Boolean = true,
     addInitialConstraint: Option[W] = None)(implicit space: MutableInnerProductModule[W, Double]) =
     {
@@ -40,34 +40,38 @@ object Trainer {
         initialConstraintAndAlpha = addInitialConstraint.map(c => ((c, 0.0), C)))
     }
 
-  def newPerceptronTrainer[T, Y, W](
-    oracleInferencer: OracleInferencer[T, Y, W], argmaxer: ArgmaxInferencer[T, Y, W],
-    iterationCallback: IterationCallback[T, W], learningRate: Double = 1.0,
-    maxNumIters: Int = 100, opts: LogoOpts = new LogoOpts(), average: Boolean = true)(implicit space: MutableInnerProductModule[W, Double]) =
+  def newPerceptronTrainer[T, Y, W, OracleS, MaxerS](
+    oracleInferencer: OracleInferencer[T, W, OracleS], argmaxer: ArgmaxInferencer[T, W, MaxerS],
+    iterationCallback: IterationCallback[T, W, OracleS, MaxerS], learningRate: Double = 1.0,
+    maxNumIters: Int = 100,
+    opts: LogoOpts = new LogoOpts(),
+    average: Boolean = true)(implicit space: MutableInnerProductModule[W, Double]) =
     {
       val decoder = new MaxMarginDecoder(oracleInferencer, argmaxer)
       val updater = new FixedStepSizeUpdater[W](_ => learningRate, Double.PositiveInfinity)
       val convergenceChecker = new FixedIterationConvergenceChecker[W](maxNumIters)
-      new Trainer(convergenceChecker, iterationCallback, decoder, updater, opts, online = true, average = average)
+      new Trainer(convergenceChecker, iterationCallback, decoder, updater, opts, online = true,
+          average = average)
     }
 
-  def newL2MaxMarginTrainer[T, Y, W](
-    oracleInferencer: OracleInferencer[T, Y, W], argmaxer: LossAugmentedArgmaxInferencer[T, Y, W],
-    iterationCallback: IterationCallback[T, W], C: Double = 1.0,
+  def newL2MaxMarginTrainer[T, Y, W, OracleS, MaxerS](
+    oracleInferencer: OracleInferencer[T, W, OracleS], argmaxer: LossAugmentedArgmaxInferencer[T, W, MaxerS],
+    iterationCallback: IterationCallback[T, W, OracleS, MaxerS], C: Double = 1.0,
     maxNumIters: Int = 100, opts: LogoOpts = new LogoOpts(),
     addInitialConstraint: Option[W] = None)(implicit space: MutableInnerProductModule[W, Double]) =
     {
       val decoder = new LossAugmentedMaxMarginDecoder(oracleInferencer, argmaxer)
       val updater = new L2Updater[W](C)
       val objective = new L2Objective[W](C)
-      val convergenceChecker = new ObjectiveFunctionConvergenceChecker(objective, maxNumIters, iterationCallback,opts.convergenceTolerance)
+      val convergenceChecker = new ObjectiveFunctionConvergenceChecker(objective, maxNumIters,
+        iterationCallback, opts.convergenceTolerance)
       new Trainer(convergenceChecker, iterationCallback, decoder, updater, opts, online = false,
         initialConstraintAndAlpha = addInitialConstraint.map(c => ((c, 0.0), C)))
     }
 
-  def newL1LogLossTrainer[T, Y, W](
-    oracleInferencer: OracleInferencer[T, Y, W], summer: ExpectationInferencer[T, W],
-    iterationCallback: IterationCallback[T, W], C: Double = 1.0,
+  def newL1LogLossTrainer[T, Y, W, OracleS, MaxerS](
+    oracleInferencer: OracleInferencer[T, W, OracleS], summer: ExpectationInferencer[T, W, MaxerS],
+    iterationCallback: IterationCallback[T, W, OracleS, MaxerS], C: Double = 1.0,
     maxNumIters: Int = 100, opts: LogoOpts = new LogoOpts(),
     addInitialConstraint: Option[W] = None)(implicit space: MutableInnerProductModule[W, Double]) =
     {
@@ -80,9 +84,9 @@ object Trainer {
         initialConstraintAndAlpha = addInitialConstraint.map(c => ((c, 0.0), C)))
     }
 
-  def newL1LogLossMIRATrainer[T, Y, W](
-    oracleInferencer: OracleInferencer[T, Y, W], summer: ExpectationInferencer[T, W],
-    iterationCallback: IterationCallback[T, W], C: Double = 1.0,
+  def newL1LogLossMIRATrainer[T, W, OracleS, MaxerS](
+    oracleInferencer: OracleInferencer[T, W, OracleS], summer: ExpectationInferencer[T, W, MaxerS],
+    iterationCallback: IterationCallback[T, W, OracleS, MaxerS], C: Double = 1.0,
     maxNumIters: Int = 100, opts: LogoOpts = new LogoOpts(), average: Boolean = true,
     addInitialConstraint: Option[W] = None)(implicit space: MutableInnerProductModule[W, Double]) =
     {
@@ -93,10 +97,10 @@ object Trainer {
           average = average, initialConstraintAndAlpha = addInitialConstraint.map(c => ((c, 0.0), C)))
     }
 
-  def newStochasticGradientDescentTrainer[T, Y, W](
-    oracleInferencer: OracleInferencer[T, Y, W],
-    summer: ExpectationInferencer[T, W],
-    iterationCallback: IterationCallback[T, W], C: Double = 1.0,
+  def newStochasticGradientDescentTrainer[T, W, OracleS, MaxerS](
+    oracleInferencer: OracleInferencer[T, W, OracleS],
+    summer: ExpectationInferencer[T, W, MaxerS],
+    iterationCallback: IterationCallback[T, W, OracleS, MaxerS], C: Double = 1.0,
     learningRate: Double = 1.0,
     maxNumIters: Int = 100, opts: LogoOpts = new LogoOpts(),
     average: Boolean = true)(implicit space: MutableInnerProductModule[W, Double]) =
@@ -108,11 +112,12 @@ object Trainer {
           average = average)
     }
 
-  def newL1MarginRankTrainer[T, Y, W](argmaxer: LossAugmentedArgmaxInferencer[T, Y, W],
-                                      iterationCallback: IterationCallback[T, W], C: Double = 1.0,
-                                      gamma: Double = 0.0, maxNumIters: Int = 100,
-                                      opts: LogoOpts = new LogoOpts(),
-                                      addInitialConstraint: Option[W] = None)(implicit space: MutableInnerProductModule[W, Double]) =
+  def newL1MarginRankTrainer[T, W, MaxerS](
+    argmaxer: LossAugmentedArgmaxInferencer[T, W, MaxerS],
+    iterationCallback: IterationCallback[T, W, MaxerS, MaxerS], C: Double = 1.0,
+    gamma: Double = 0.0, maxNumIters: Int = 100,
+    opts: LogoOpts = new LogoOpts(),
+    addInitialConstraint: Option[W] = None)(implicit space: MutableInnerProductModule[W, Double]) =
     {
       val decoder = new MaxMarginRankingDecoder(argmaxer, gamma)
       val updater = new L1Updater[W](C)
@@ -128,8 +133,8 @@ object Trainer {
     data: Seq[LabeledDatum[L, F]],
     labelConjoiner: (L, F) => W,
     initialConstraint: W,
-    iterationCallback: IterationCallback[LabeledDatum[L, F], W] =
-      new NullIterationCallback[LabeledDatum[L, F], W](),
+    iterationCallback: IterationCallback[LabeledDatum[L, F], W, Unit, Unit] =
+      new NullIterationCallback[LabeledDatum[L, F], W, Unit, Unit](),
     oneSlackFormulation: Boolean = true,
     C: Double = 1.0,
     maxNumIters: Int = 100, opts: LogoOpts = new LogoOpts())(implicit space: MutableInnerProductModule[W, Double]) = {
@@ -137,12 +142,12 @@ object Trainer {
     val argmaxer = new MulticlassLossAugmentedArgmaxInferencer[L, F, W](labels, labelConjoiner)
     val oracler = new MulticlassOracleInferencer[L, F, W](labels, labelConjoiner)
     val weights = if (oneSlackFormulation) {
-      val oneSlackIterCallBack = new IterationCallback[Seq[LabeledDatum[L, F]], W]() {
+      val oneSlackIterCallBack = new IterationCallback[Seq[LabeledDatum[L, F]], W, Unit, Unit]() {
         override def startIteration(iter: Int, weights: Weights[W]): Unit =
           iterationCallback.startIteration(iter, weights)
 
-        override def endIteration(iter: Int, weights: Weights[W]): Unit =
-          iterationCallback.endIteration(iter, weights)
+        override def endIteration(iter: Int, weights: Weights[W], unused: Unit, unused2: Unit): Unit =
+          iterationCallback.endIteration(iter, weights, unused, unused2)
 
         override def objectiveValCheck(primal: Double, dual: Double, iter: Int, weights: Weights[W]): Unit =
           iterationCallback.objectiveValCheck(primal, dual, iter, weights)
@@ -166,13 +171,15 @@ object Trainer {
   }
 }
 
-case class Trainer[T, W](convergenceChecker: ConvergenceChecker[W],
-                         iterationCallback: IterationCallback[T, W], val decoder: Decoder[T, W],
-                         updater: Updater[W],
-                         opts: LogoOpts = new LogoOpts(),
-                         initialConstraintAndAlpha: Option[((W, Double), Double)] = None,
-                         online: Boolean = false,
-                         average: Boolean = false)(implicit space: MutableInnerProductModule[W, Double])
+case class Trainer[T, W, OracleS, MaxerS](
+  convergenceChecker: ConvergenceChecker[W],
+  iterationCallback: IterationCallback[T, W, OracleS, MaxerS],
+  decoder: Decoder[T, W, OracleS, MaxerS],
+  updater: Updater[W],
+  opts: LogoOpts = new LogoOpts(),
+  initialConstraintAndAlpha: Option[((W, Double), Double)] = None,
+  online: Boolean = false,
+  average: Boolean = false)(implicit space: MutableInnerProductModule[W, Double])
     extends SerializableLogging {
   import space._
 
@@ -205,6 +212,7 @@ case class Trainer[T, W](convergenceChecker: ConvergenceChecker[W],
       iterationCallback.startIteration(iteration, w)
 
       var currStart = 0;
+      var (oracleInferencerState, maxerInferenceState) = decoder.initialState
       while (iter.hasNext) {
         val currMiniBatchUnshuffled = consume(iter, Math.min(data.length, opts.miniBatchSize))
         val currMiniBatch =
@@ -214,7 +222,7 @@ case class Trainer[T, W](convergenceChecker: ConvergenceChecker[W],
           currMiniBatchUnshuffled
         iterationCallback.startMinibatch(iteration, w, currMiniBatch)
         val decodedMiniBatch = for (MinibatchInput(instance, instanceNum) <- currMiniBatch) yield {
-          val (df, l) = decoder.decode(w, instance.x)
+          val (df, l, newOracleInferencerState, newMaxerInferenceState) = decoder.decode(w, instance.x)
           if (!online) {
             if (iteration == 0 && initialConstraintAndAlpha.isDefined) {
               val ((df_, l_), alpha) = initialConstraintAndAlpha.get
@@ -234,12 +242,19 @@ case class Trainer[T, W](convergenceChecker: ConvergenceChecker[W],
               numAdded += 1
             }
           }
-          MinibatchOutput(instance, instanceNum, df, l)
+          MinibatchOutput(instance, instanceNum, df, l, newOracleInferencerState, newMaxerInferenceState)
+        }
+        val (newOracleInferencerState, newMaxerInferenceState) =
+          decodedMiniBatch.foldLeft((oracleInferencerState, maxerInferenceState)) {
+          case (statePair, MinibatchOutput(_, _, _, _, oracleState, maxerState)) =>
+            decoder.reduceStates(statePair, (oracleState, maxerState))
         }
+        oracleInferencerState = newOracleInferencerState
+        maxerInferenceState = newMaxerInferenceState
         iterationCallback.endMinibatch(iteration, w, decodedMiniBatch)
         loopWhile(numInnerOptimizationLoops) {
           val numChanges = decodedMiniBatch.count {
-            case MinibatchOutput(instance, instanceNum, df, l) =>
+            case MinibatchOutput(instance, instanceNum, df, l, _, _) =>
               if (online) {
                 val constraints = ArrayBuffer((df, l)) ++ initialConstraintAndAlpha.map(_._1).toList
                 val alphas = ArrayBuffer(0.0) ++ initialConstraintAndAlpha.map(_._2).toList
@@ -262,7 +277,7 @@ case class Trainer[T, W](convergenceChecker: ConvergenceChecker[W],
         }
         instancesChanged > 0
       }
-      iterationCallback.endIteration(iteration, w)
+      iterationCallback.endIteration(iteration, w, oracleInferencerState, maxerInferenceState)
       if (average) {
         average_w *= iteration / (iteration + 1)
         average_w.increment(w, 1.0 / (iteration + 1))

From bdf0547a0919d604cf236bdc542181e413f7e9a8 Mon Sep 17 00:00:00 2001
From: adpauls <adpauls@semanticmachines.com>
Date: Thu, 26 May 2016 19:12:53 -0700
Subject: [PATCH 6/7] Whoops.

---
 src/main/scala/epic/logo/Trainer.scala | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/main/scala/epic/logo/Trainer.scala b/src/main/scala/epic/logo/Trainer.scala
index 6b93a500..a4ac7115 100644
--- a/src/main/scala/epic/logo/Trainer.scala
+++ b/src/main/scala/epic/logo/Trainer.scala
@@ -247,6 +247,7 @@ case class Trainer[T, W, OracleS, MaxerS](
         val (newOracleInferencerState, newMaxerInferenceState) =
           decodedMiniBatch.foldLeft((oracleInferencerState, maxerInferenceState)) {
           case (statePair, MinibatchOutput(_, _, _, _, oracleState, maxerState)) =>
+
             decoder.reduceStates(statePair, (oracleState, maxerState))
         }
         oracleInferencerState = newOracleInferencerState
@@ -268,14 +269,14 @@ case class Trainer[T, W, OracleS, MaxerS](
       loopWhile(numOuterOptimizationLoops) {
         // Note: we don't use exists because we want to run on every example, not stop on the first
         // one with a change
-        val instancesChanged = data.count { instance =>
+        val numChanges = data.count { instance =>
           val numUpdatesExecuted = loopWhile(numInnerOptimizationLoops) {
             updater.update(instance, w, n, iteration)
           }
           // 1 instead of 0 because update is always called once, but it might do nothing.
-          numUpdatesExecuted <= 1
+          numUpdatesExecuted > 1
         }
-        instancesChanged > 0
+        numChanges > 0
       }
       iterationCallback.endIteration(iteration, w, oracleInferencerState, maxerInferenceState)
       if (average) {

From 5a33ccfdaeec5214bb190cf6d4d3856b2e4a0e86 Mon Sep 17 00:00:00 2001
From: adpauls <adpauls@semanticmachines.com>
Date: Thu, 26 May 2016 22:21:47 -0700
Subject: [PATCH 7/7] PR comments

---
 .../scala/epic/logo/CompoundIterationCallback.scala  |  2 +-
 src/main/scala/epic/logo/ConvergenceChecker.scala    |  2 +-
 src/main/scala/epic/logo/DualVariableHolder.scala    | 12 ++++++++++++
 .../epic/logo/FixedIterationConvergenceChecker.scala |  2 +-
 src/main/scala/epic/logo/FixedStepSizeUpdater.scala  |  4 ++--
 src/main/scala/epic/logo/Inferencer.scala            |  4 ++++
 src/main/scala/epic/logo/Instance.scala              | 11 -----------
 src/main/scala/epic/logo/IterationCallback.scala     |  2 +-
 src/main/scala/epic/logo/L1Objective.scala           |  4 ++--
 src/main/scala/epic/logo/L1Updater.scala             |  4 ++--
 src/main/scala/epic/logo/L2Objective.scala           |  2 +-
 src/main/scala/epic/logo/L2Updater.scala             |  4 ++--
 src/main/scala/epic/logo/MinibatchInput.scala        |  4 ++--
 src/main/scala/epic/logo/ObjectiveFunction.scala     |  2 +-
 .../logo/ObjectiveFunctionConvergenceChecker.scala   |  4 ++--
 src/main/scala/epic/logo/PegasosUpdater.scala        |  4 ++--
 src/main/scala/epic/logo/Trainer.scala               |  8 ++++----
 src/main/scala/epic/logo/Updater.scala               |  4 ++--
 18 files changed, 42 insertions(+), 37 deletions(-)
 create mode 100644 src/main/scala/epic/logo/DualVariableHolder.scala
 delete mode 100644 src/main/scala/epic/logo/Instance.scala

diff --git a/src/main/scala/epic/logo/CompoundIterationCallback.scala b/src/main/scala/epic/logo/CompoundIterationCallback.scala
index aaa746f5..7099776d 100644
--- a/src/main/scala/epic/logo/CompoundIterationCallback.scala
+++ b/src/main/scala/epic/logo/CompoundIterationCallback.scala
@@ -24,7 +24,7 @@ class CompoundIterationCallback[T, W, S1, S2](val callbacks: Iterable[IterationC
     callbacks.foreach(_.objectiveValCheck(primal, dual, iter, w))
   }
 
-  override def converged(weights: Weights[W], data: Seq[Instance[T, W]],
+  override def converged(weights: Weights[W], data: Seq[DualVariableHolder[T, W]],
                          iter: Int, numNewConstraints: Int): Boolean = {
     callbacks.forall(_.converged(weights, data, iter, numNewConstraints))
   }
diff --git a/src/main/scala/epic/logo/ConvergenceChecker.scala b/src/main/scala/epic/logo/ConvergenceChecker.scala
index bc642b4e..b76d8a89 100644
--- a/src/main/scala/epic/logo/ConvergenceChecker.scala
+++ b/src/main/scala/epic/logo/ConvergenceChecker.scala
@@ -2,6 +2,6 @@ package epic.logo
 
 trait ConvergenceChecker[W] {
 
-  def converged(weights : Weights[W], data : Seq[Instance[_, W]], iter : Int, numNewConstraints : Int) : Boolean
+  def converged(weights : Weights[W], data : Seq[DualVariableHolder[_, W]], iter : Int, numNewConstraints : Int) : Boolean
 
 }
diff --git a/src/main/scala/epic/logo/DualVariableHolder.scala b/src/main/scala/epic/logo/DualVariableHolder.scala
new file mode 100644
index 00000000..f3a4d2c8
--- /dev/null
+++ b/src/main/scala/epic/logo/DualVariableHolder.scala
@@ -0,0 +1,12 @@
+package epic.logo
+
+import scala.collection.mutable.ArrayBuffer
+
+/**
+  * This class stores the dual variables and constraints for each training exampe.
+  * It is mutable for efficiency purposes.
+  */
+case class DualVariableHolder[T, W](var x: T,
+                                    var alphas: ArrayBuffer[Double] = ArrayBuffer[Double](),
+                                    var slack: Double = 0.0,
+                                    var constraints: ArrayBuffer[(W, Double)] = ArrayBuffer[(W, Double)]())
diff --git a/src/main/scala/epic/logo/FixedIterationConvergenceChecker.scala b/src/main/scala/epic/logo/FixedIterationConvergenceChecker.scala
index b8303cc3..e4f5a479 100644
--- a/src/main/scala/epic/logo/FixedIterationConvergenceChecker.scala
+++ b/src/main/scala/epic/logo/FixedIterationConvergenceChecker.scala
@@ -4,7 +4,7 @@ import scala.collection.Seq
 
 class FixedIterationConvergenceChecker[W](val maxNumIters : Int) extends ConvergenceChecker[W] {
 
-  def converged(weights : Weights[W], data : Seq[Instance[_, W]], iter : Int, numNewConstraints : Int) : Boolean = {
+  def converged(weights : Weights[W], data : Seq[DualVariableHolder[_, W]], iter : Int, numNewConstraints : Int) : Boolean = {
     iter >= maxNumIters
   }
 
diff --git a/src/main/scala/epic/logo/FixedStepSizeUpdater.scala b/src/main/scala/epic/logo/FixedStepSizeUpdater.scala
index da644e33..ba76af4d 100644
--- a/src/main/scala/epic/logo/FixedStepSizeUpdater.scala
+++ b/src/main/scala/epic/logo/FixedStepSizeUpdater.scala
@@ -7,7 +7,7 @@ import breeze.math.MutableInnerProductModule
 class FixedStepSizeUpdater[W](stepSize : Int => Double, C : Double)(implicit space: MutableInnerProductModule[W, Double]) extends Updater[W] {
   import space._
 
-  def update(instance: Instance[_, W],
+  def update(instance: DualVariableHolder[_, W],
              w: Weights[W], n: Int, iter: Int): Boolean = {
     import instance._
     assert(constraints.length == 2)
@@ -24,7 +24,7 @@ class FixedStepSizeUpdater[W](stepSize : Int => Double, C : Double)(implicit spa
     return true
   }
 
-  def currentSlack(i : Instance[_, W], w : Weights[W]) : Double = {
+  def currentSlack(i : DualVariableHolder[_, W], w : Weights[W]) : Double = {
     throw new UnsupportedOperationException(this.getClass().getName() + " should be only be used in online mode.")
   }
 
diff --git a/src/main/scala/epic/logo/Inferencer.scala b/src/main/scala/epic/logo/Inferencer.scala
index 5ad680be..2ac2fa64 100644
--- a/src/main/scala/epic/logo/Inferencer.scala
+++ b/src/main/scala/epic/logo/Inferencer.scala
@@ -1,6 +1,10 @@
 package epic.logo
 
 trait Inferencer[S] {
+  /** Inferencers may wish to carry state between calls. Rather than maintaining it internally
+   *  and breaking referentially transparency, the may return states after each call that
+   *  are aggregated together using [[reduceStates]].
+   */
   def initialState: S
   def reduceStates(state1: S, state2: S): S
 
diff --git a/src/main/scala/epic/logo/Instance.scala b/src/main/scala/epic/logo/Instance.scala
deleted file mode 100644
index 7d11321f..00000000
--- a/src/main/scala/epic/logo/Instance.scala
+++ /dev/null
@@ -1,11 +0,0 @@
-package epic.logo
-
-import scala.collection.mutable.ArrayBuffer
-
-/**
- * This class is mutable for efficiency purposes.
- */
-case class Instance[T, W](var x: T,
-                          var alphas: ArrayBuffer[Double] = ArrayBuffer[Double](),
-                          var slack: Double = 0.0,
-                          var constraints: ArrayBuffer[(W, Double)] = ArrayBuffer[(W, Double)]())
diff --git a/src/main/scala/epic/logo/IterationCallback.scala b/src/main/scala/epic/logo/IterationCallback.scala
index 1185a25e..07dbed41 100644
--- a/src/main/scala/epic/logo/IterationCallback.scala
+++ b/src/main/scala/epic/logo/IterationCallback.scala
@@ -10,7 +10,7 @@ trait IterationCallback[T, W, OracleS, MaxerS] {
 
   def objectiveValCheck(primal : Double, dual : Double, iter : Int, weights : Weights[W]) : Unit = {}
 
-  def converged(weights : Weights[W], data : Seq[Instance[T, W]], iter : Int, numNewConstraints : Int) : Boolean = false
+  def converged(weights : Weights[W], data : Seq[DualVariableHolder[T, W]], iter : Int, numNewConstraints : Int) : Boolean = false
 
 }
 
diff --git a/src/main/scala/epic/logo/L1Objective.scala b/src/main/scala/epic/logo/L1Objective.scala
index 98476503..3dfa13c6 100644
--- a/src/main/scala/epic/logo/L1Objective.scala
+++ b/src/main/scala/epic/logo/L1Objective.scala
@@ -4,7 +4,7 @@ import breeze.linalg._
 import breeze.math.MutableInnerProductModule
 
 object L1Objective {
-  def slack[W](i : Instance[_, W], w : Weights[W])(implicit space: MutableInnerProductModule[W, Double]) = {
+  def slack[W](i : DualVariableHolder[_, W], w : Weights[W])(implicit space: MutableInnerProductModule[W, Double]) = {
     import space._
     if (i.constraints.isEmpty)
       Double.NegativeInfinity
@@ -17,7 +17,7 @@ object L1Objective {
 class L1Objective[W](val C : Double)(implicit space: MutableInnerProductModule[W, Double]) extends ObjectiveFunction[W] {
   import space._
 
-  def calculatePrimalAndDual(w : Weights[W], data : Seq[Instance[_, W]]): (Double, Double) = {
+  def calculatePrimalAndDual(w : Weights[W], data : Seq[DualVariableHolder[_, W]]): (Double, Double) = {
     w.checkNorm
     val calc_w = new Weights(space.zeroLike(w.underlying))
     data.flatMap(instance => instance.alphas zip instance.constraints).map { case (alpha, (df, l)) => df * alpha }.foreach(fv => {
diff --git a/src/main/scala/epic/logo/L1Updater.scala b/src/main/scala/epic/logo/L1Updater.scala
index edcf9ebe..02cd605c 100644
--- a/src/main/scala/epic/logo/L1Updater.scala
+++ b/src/main/scala/epic/logo/L1Updater.scala
@@ -9,7 +9,7 @@ class L1Updater[W](C : Double)(implicit space: MutableInnerProductModule[W, Doub
   val shuffleRand = new Random(1)
   import space._
 
-  def update(instance: Instance[_, W], w : Weights[W], n : Int, iter : Int) : Boolean = {
+  def update(instance: DualVariableHolder[_, W], w : Weights[W], n : Int, iter : Int) : Boolean = {
     import instance._
     if (constraints.length == 1) {
       val alpha0 = alphas(0)
@@ -49,7 +49,7 @@ class L1Updater[W](C : Double)(implicit space: MutableInnerProductModule[W, Doub
     }
   }
 
-  def currentSlack(i : Instance[_, W], w : Weights[W]) = {
+  def currentSlack(i : DualVariableHolder[_, W], w : Weights[W]) = {
     L1Objective.slack(i, w)
   }
 }
diff --git a/src/main/scala/epic/logo/L2Objective.scala b/src/main/scala/epic/logo/L2Objective.scala
index aadc28d1..c3efbc74 100644
--- a/src/main/scala/epic/logo/L2Objective.scala
+++ b/src/main/scala/epic/logo/L2Objective.scala
@@ -5,7 +5,7 @@ import breeze.math.MutableInnerProductModule
 class L2Objective[W](val C: Double)(implicit space: MutableInnerProductModule[W, Double])
     extends ObjectiveFunction[W] {
 
-  def calculatePrimalAndDual(w: Weights[W], data: Seq[Instance[_, W]]): (Double, Double) = {
+  def calculatePrimalAndDual(w: Weights[W], data: Seq[DualVariableHolder[_, W]]): (Double, Double) = {
     val primal = {
       val slackSum = data.map(instance => instance.slack * instance.slack).sum
       0.5 * (w.`^2`) + C * slackSum
diff --git a/src/main/scala/epic/logo/L2Updater.scala b/src/main/scala/epic/logo/L2Updater.scala
index ca8a5dc2..088b714a 100644
--- a/src/main/scala/epic/logo/L2Updater.scala
+++ b/src/main/scala/epic/logo/L2Updater.scala
@@ -7,7 +7,7 @@ import breeze.math.MutableInnerProductModule
 class L2Updater[W](C : Double)(implicit space: MutableInnerProductModule[W, Double]) extends Updater[W] {
   import space._
 
-  def update(instance: Instance[_, W],
+  def update(instance: DualVariableHolder[_, W],
              w: Weights[W], n: Int, iter: Int): Boolean = {
     import instance._
     for (((df, l), s) <- constraints.zipWithIndex) {
@@ -31,7 +31,7 @@ class L2Updater[W](C : Double)(implicit space: MutableInnerProductModule[W, Doub
 
   }
 
-  def currentSlack(i : Instance[_, W], w : Weights[W]) = {
+  def currentSlack(i : DualVariableHolder[_, W], w : Weights[W]) = {
     i.slack
   }
 
diff --git a/src/main/scala/epic/logo/MinibatchInput.scala b/src/main/scala/epic/logo/MinibatchInput.scala
index d6456012..59d1b5e3 100644
--- a/src/main/scala/epic/logo/MinibatchInput.scala
+++ b/src/main/scala/epic/logo/MinibatchInput.scala
@@ -1,7 +1,7 @@
 package epic.logo
 
-case class MinibatchInput[T, W](instance : Instance[T, W], instanceNum : Int)
+case class MinibatchInput[T, W](instance : DualVariableHolder[T, W], instanceNum : Int)
 
 case class MinibatchOutput[T, W, OracleS, MaxerS](
-  instance: Instance[T, W], instanceNum: Int, df: W, loss: Double,
+  instance: DualVariableHolder[T, W], instanceNum: Int, df: W, loss: Double,
   oracleState: OracleS, maxerState: MaxerS)
diff --git a/src/main/scala/epic/logo/ObjectiveFunction.scala b/src/main/scala/epic/logo/ObjectiveFunction.scala
index 5549aa1f..45d85181 100644
--- a/src/main/scala/epic/logo/ObjectiveFunction.scala
+++ b/src/main/scala/epic/logo/ObjectiveFunction.scala
@@ -2,6 +2,6 @@ package epic.logo
 
 trait ObjectiveFunction[W] {
 
-  def calculatePrimalAndDual(w : Weights[W], data : Seq[Instance[_, W]]) : (Double, Double)
+  def calculatePrimalAndDual(w : Weights[W], data : Seq[DualVariableHolder[_, W]]) : (Double, Double)
 
 }
diff --git a/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala b/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala
index cddc2f15..83f63715 100644
--- a/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala
+++ b/src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala
@@ -6,11 +6,11 @@ case class ObjectiveFunctionConvergenceChecker[W](
     objective: ObjectiveFunction[W], maxNumIters: Int,
     callback: IterationCallback[_, W, _, _], tol: Double) extends ConvergenceChecker[W] {
 
-  def converged(weights: Weights[W], data: Seq[Instance[_, W]], iter: Int, numNewConstraints: Int): Boolean = {
+  def converged(weights: Weights[W], data: Seq[DualVariableHolder[_, W]], iter: Int, numNewConstraints: Int): Boolean = {
     iter >= maxNumIters || (objectiveConverged(weights, data, iter) && numNewConstraints == 0)
   }
 
-  private def objectiveConverged(weights: Weights[W], data: Seq[Instance[_, W]], iter: Int) = {
+  private def objectiveConverged(weights: Weights[W], data: Seq[DualVariableHolder[_, W]], iter: Int) = {
     val (primal, dual) = objective.calculatePrimalAndDual(weights, data)
     callback.objectiveValCheck(primal, dual, iter, weights)
     NumUtils.approxEquals(primal, dual, tol)
diff --git a/src/main/scala/epic/logo/PegasosUpdater.scala b/src/main/scala/epic/logo/PegasosUpdater.scala
index f6b5900b..431b7bd4 100644
--- a/src/main/scala/epic/logo/PegasosUpdater.scala
+++ b/src/main/scala/epic/logo/PegasosUpdater.scala
@@ -6,7 +6,7 @@ import breeze.math.MutableInnerProductModule
 
 class PegasosUpdater[W](C : Double)(implicit space: MutableInnerProductModule[W, Double]) extends Updater[W] {
   import space._
-  def update(instance: Instance[_, W], w : Weights[W], n : Int, iter : Int) : Boolean = {
+  def update(instance: DualVariableHolder[_, W], w : Weights[W], n : Int, iter : Int) : Boolean = {
     import instance._
     assert(constraints.length == 2)
     val (df, _) = constraints(0)
@@ -24,7 +24,7 @@ class PegasosUpdater[W](C : Double)(implicit space: MutableInnerProductModule[W,
 
   }
 
-  def currentSlack(i : Instance[_, W], w : Weights[W]) : Double = {
+  def currentSlack(i : DualVariableHolder[_, W], w : Weights[W]) : Double = {
     throw new UnsupportedOperationException(this.getClass().getName() + " should be only be used in online mode.")
   }
 
diff --git a/src/main/scala/epic/logo/Trainer.scala b/src/main/scala/epic/logo/Trainer.scala
index a4ac7115..6bea0d05 100644
--- a/src/main/scala/epic/logo/Trainer.scala
+++ b/src/main/scala/epic/logo/Trainer.scala
@@ -152,9 +152,9 @@ object Trainer {
         override def objectiveValCheck(primal: Double, dual: Double, iter: Int, weights: Weights[W]): Unit =
           iterationCallback.objectiveValCheck(primal, dual, iter, weights)
 
-        override def converged(weights: Weights[W], data: Seq[Instance[Seq[LabeledDatum[L, F]], W]],
+        override def converged(weights: Weights[W], data: Seq[DualVariableHolder[Seq[LabeledDatum[L, F]], W]],
                                iter: Int, numNewConstraints: Int): Boolean = {
-          iterationCallback.converged(weights, data.head.x.map(new Instance[LabeledDatum[L, F], W](_)),
+          iterationCallback.converged(weights, data.head.x.map(DualVariableHolder[LabeledDatum[L, F], W](_)),
             iter, numNewConstraints)
         }
 
@@ -188,7 +188,7 @@ case class Trainer[T, W, OracleS, MaxerS](
   final val numOuterOptimizationLoops = if (online) 0 else opts.numOuterOptimizationLoops
 
   def train(rawData : Seq[T], initWeights : Weights[W]) = {
-    val data = rawData.map(datum => Instance[T, W](datum))
+    val data = rawData.map(datum => DualVariableHolder[T, W](datum))
     val n = data.length
     val w = if (online) initWeights else {
       if (initWeights.norm != 0.0) {
@@ -259,7 +259,7 @@ case class Trainer[T, W, OracleS, MaxerS](
               if (online) {
                 val constraints = ArrayBuffer((df, l)) ++ initialConstraintAndAlpha.map(_._1).toList
                 val alphas = ArrayBuffer(0.0) ++ initialConstraintAndAlpha.map(_._2).toList
-                updater.update(Instance(constraints, alphas), w, n, iteration)
+                updater.update(DualVariableHolder(constraints, alphas), w, n, iteration)
               } else
                 updater.update(instance, w, n, iteration)
           }
diff --git a/src/main/scala/epic/logo/Updater.scala b/src/main/scala/epic/logo/Updater.scala
index 9afaef44..8db2ca55 100644
--- a/src/main/scala/epic/logo/Updater.scala
+++ b/src/main/scala/epic/logo/Updater.scala
@@ -8,9 +8,9 @@ trait Updater[W] {
   /**
    * Mutates `instance` and `weights`.
    */
-  def update(instance: Instance[_, W],
+  def update(instance: DualVariableHolder[_, W],
              weights: Weights[W], n: Int, iter: Int): Boolean
 
-  def currentSlack(i: Instance[_, W], w: Weights[W]): Double
+  def currentSlack(i: DualVariableHolder[_, W], w: Weights[W]): Double
 
 }