神经网络与函数式编程（七）如何实现Monadic Deep Learning？

神经网络与函数式编程

在本系列先前的文章神经网络与函数式编程（四）Monadic Deep Learning中，我们学习了DeepLearning.scala如何实现多态函数，把各种不同的参数统一转为Do[Tape]来处理。在本篇文章中，我们将讨论DeepLearning.scala的反向传播机制，看看DeepLearning.scala如何实现Monadic风格的可差分计算图。

方案一：基于Tape的反向传播

首先我将简化一下DeepLearning.scala的实现，把涉及Monad的部分去掉，展示一下如何DeepLearning.scala内部的反向传播流程。

以一元线性回归模型为例，如果采用LAD惩罚函数，可以用DeepLearning.scala表达成这样：

val weight = DoubleWeight(math.random)
def f(x: Double): DoubleLayer = x * weight
def lossFunction(x: Double, expectedY: Double): DoubleLayer = {
  abs(expectedY - f(x))
}

给定带标记的数据样本x和expectedY，可以用train函数训练：

val x: Double = ???
val expectedY: Double = ???
lossFunction(x, expectedY).train

其中，lossFunction(x, expectedY)是一个计算图DoubleLayer。在本系列前一篇文章，我们已经知道DoubleLayer可以forward转为Do[Tape[Double, Double]]。Do是个Monad类型，我们暂时忽略，把DoubleLayer.forward简单的看成Tape[Double, Double]。简化后的Tape定义如下：

case class Tape[Data, Delta](data: Data, backward: Delta => Unit)

Tape包含前向阶段的结果data，以及用来进行反向传播的backward函数。

比如train是反向传播的起点，可以实现成这样：

trait DeepLearning[Differentiable] {
  type Data
  type Delta
  def forward(differentiable: Differentiable): Tape[Data, Delta]
  def train(lossFunction: DoubleLayer): Data = {
    val tape: Tape[Data, Delta] = lossFunction.forward
    tape.backward(1.0)
    tape.data
  }
}

train用到的lossFunction展开后是abs(expectedY - x * weight)。

// TODO: 此处缺一张图

其中可差分的abs实现如下

def abs(operand0: DoubleLayer): DoubleLayer = {
  def forward: Tape[Double, Double] = {
    val tape = operand0.forward
    val data = math.abs(tape.data)
    val backward = { delta =>
      tape.backward(math.sign(tape.data) * delta)
    }
    Tape(data, backward)
  }
  DoubleLayer(forward)
}

可以看到在前向阶段，abs的data是operand0的data的绝对值。而在反向传播时，delta函数把导数经过计算后，传给operand0的backward。

而abs用到的operand0则是可差分的减法，根据求导公式可以写成这样：

def -(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
  def forward: Tape[Double, Double] = {
    val tape0 = operand0.forward
    val tape1 = operand1.forward
    val data = tape0.data - tape1.data
    val backward = { delta =>
      tape0.backward(delta)
      tape1.backward(-delta)
    }
    Tape(data, backward)
  }
  DoubleLayer(forward)
}

注意，减法与abs不同，是个二元运算符，所以backward时需要分别反向传播到两个操作数的backward。

然后，减法的operand1是f(x)，是个可差分的乘法：

def *(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
  def forward: Tape[Double, Double] = {
    val tape0 = operand0.forward
    val tape1 = operand1.forward
    val data = tape0.data * tape1.data
    val backward = { delta =>
      tape0.backward(delta * tape1.data)
      tape1.backward(delta * tape0.data)
    }
    Tape(data, backward)
  }
  DoubleLayer(forward)
}

最后，weight是个DoubleWeight。DoubleWeight的backward会修改自己内部保存的值：

trait DoubleWeight {
  var data: Double
  def forward: Tape[Double, Double] = {
    val backward = { delta =>
      data -= delta
    }
    Tape(data, backward)
  }
}

到此为止，可以看到，当用户调用train时。
train内部首先会递归调用forward，在最深处由DoubleWeight取得data，然后一层一层执行前向阶段，返回一个Tape，其中包含前些阶段的执行结果。
然后，train内部会递归调用Tape上的backward，把导数反向传播，直到传到DoubleWeight上，令DoubleWeight得到训练。

方案二：支持粗粒度并行计算的异步计算图

到目前为止，我们实现的计算图是同步执行的。这意味着，计算图上的节点都不可能并行执行，即使这些节点之间并没有相互依赖关系。为了实现神经网络与函数式编程（四）Monadic Deep Learning中提到的粗粒度并行计算，我们需要让forward返回一个异步任务，而不应该阻塞。

trait DeepLearning[Differentiable] {
  type Data
  type Delta
  def forward(differentiable: Differentiable): UnitContinuation[Tape[Data, Delta]]
  def train(lossFunction: DoubleLayer): UnitContinuation[Data] = monadic[UnitContinuation] {
    val tape = lossFunction.forward.each
    tape.backward(UnitContinuation.now(1.0)).each
    tape.data
  }
}

以上代码中，我们引入了UnitContinuation来表示异步任务。train也做了修改，改用monadic风格的控制流。

相应的Tape.backward也需要改为异步结构。

case class Tape[Data, Delta](data: Data, backward: UnitContinuation[Delta] => UnitContinuation[Unit])

我们用了一个小技巧，让backward接受UnitContinuation[Delta]而不是Delta参数。这样一来，由于UnitContinuation[Delta]是个异步操作而没有严格求值，那么，当backward触发时，backward的实现者如果是不包含任何权重的计算图，backward就可以跳过不执行，最终以自动实现了类似PyTorch中requires_grad的优化效果，而无需用户任何手动设置。

最后，每一个原子操作也改为monadic风格，比如乘法操作：

def *(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
  def forward: UnitContinuation[Tape[Double, Double]] = {
    val Parallel(result) = (Parallel(operand0.forward) |@| Parallel(operand1.forward)) { (tape0, tape1) =>
      val data = tape0.data * tape1.data
      val backward = { deltaContinuation =>
        deltaContinuation.flatMap { delta =>
          val backward0 = tape0.backward(delta * tape1.data)
          val backward1 = tape1.backward(delta * tape0.data)
          val Parallel(backwardBoth) = Parallel(backward0) |@| Parallel(backward1) { (_, _) => }
          backwardBoth
        }
      }
      Tape(data, backward)
    }
    result
  }
  DoubleLayer(forward)
}

注意，我们实现了乘法这样的二元操作符时，无论是前向阶段还是反向阶段，都利用Parallel让两个操作数上的计算得到并行执行。

到此为止，方案二通过引入UnitContinuation，让所有计算图都得到了异步计算的能力，并且可以并行执行多个计算图的分支。

问题：菱形计算图导致指数时间复杂度

与autograd等早期的自动求导算法不同，DeepLearning.scala在二元操作处，能对两个操作数都求导。你可以在减法和乘法的上看到，backward内部会分别触发两个操作数的backward。这种做法可以支持复杂的网络结构.

不过，当计算图中出现菱形依赖时，要是直接在backward中反向传播就会导致指数级增长的运算量。考虑以下计算图：

val weight: DoubleLayer = DoubleWeight(math.random)
def diamondComputationalGraph: DoubleLayer = {
  val square0: DoubleLayer = weight * weight
  val square1: DoubleLayer = square0 * square0
  val square2: DoubleLayer = square1 * square1
  val square3: DoubleLayer = square2 * square2
  square3
}

当调用diamondComputationalGraph.train训练时，在反向传播阶段，首先触发square3的backward，然后square3会调用左右操作数的backward进行反向传播。由于square3的左右操作数都是square2，所以square2的backward会触发两次。而每次square2的backward会触发两次square1的backward。square1的backward一共就触发了4次。以此类推，最终weight的backward会触发16次。

方案三：基于引用计数的Resource

为了解决这个问题，我们增加了一个引用计数功能，我们让forward返回的不是Tape，而是带有引用计数的Resource：

case class Resource[A](value: A, release: UnitContinuation[Unit])

那么现在的乘法运算符可以实现成这样：

def *(operand0: DoubleLayer, operand1: DoubleLayer): DoubleLayer = {
  var referenceCount = 0
  var deltaAccumulator = 0.0
  def forward: UnitContinuation[Resource[Tape[Double, Double]]] = monadic[UnitContinuation] {
    referenceCount += 1
    val resource0 = operand0.forward.each
    val resource1 = operand1.forward.each
    val data = resource0.value.data * resource1.value.data
    val backward = { delta =>
      monadic[UnitContinuation] {
        deltaAccumulator += delta.each
      }
    }
    val tape = Tape(data, backward)
    val release = monadic[UnitContinuation] {
      referenceCount -= 1
      if (referenceCount == 0) {
        resource0.value.backward(deltaAccumulator * resource1.value.data)
        resource1.value.backward(deltaAccumulator * resource0.value.data)
        deltaAccumulator = 0.0
        resource0.release().each
        resource1.release().each
      }
    }
    Resource(tape, release)
  }
  DoubleLayer(forward)
}

相比第二版的乘法，我们做了如下修改：
* 每个DoubleLayer会记录forward的次数。
* backward不再立即触发递归的反向传播，而是只把delta暂存到deltaAccumulator。
* 增加了release函数，release函数会释放引用计数，当引用计数释放到零时，触发反向传播并且释放上游的引用计数。

train也需要修改，在发起反向传播之处时，得在backward之后调用一次release：

def train(lossFunction: DoubleLayer): UnitContinuation[Data] = monadic[UnitContinuation] {
  val resource: Resource[Tape[Data, Delta]] = lossFunction.forward.each
  val data = resource.data
  resource.value.backward(1.0).each
  resource.release().each
}

经过这样的修改，现在diamondComputationalGraph.train训练时，只调用了一次lossFunction.forward，所以resource的引用计数为1。

然后，当train调用resource.value.backward时，新的backward实现只会把delta暂存到deltaAccumulator，不触发递归的反向传播。

最后，当train调用release时，引用计数减为零，触发递归的反向传播。

当反向传播到square3时，由于square2既是square3的左操作数也是右操作数，所以square2会被square3引用两次，引用计数为2。当square3触发release时，会同时触发square3的左操作数和右操作数的release，也就导致square2两遍，刚好可以把引用计数减为零，继续触发square2的反向传播。

以此类推，计算图的每个节点都有引用计数。即使在菱形计算图的情况下，也可以保证只在最后一次release时才触发反向传播，而不会导致指数增长的运算量。

问题：动态神经网络的引用计数

这种基于引用计数的算法，必须保证forward调用次数和release相同。然而，在复杂的动态网络中很难保证。比如：

val weight: DoubleLayer = DoubleWeight(math.random)
def diamondComputationalGraph: DoubleLayer = DoubleLayer(monadic[UnitContinuation] {
  if (weight.forward.each.value.data > 0.5) {
    val square0: DoubleLayer = weight * weight
    if (square0.forward.each.value.data > 0.5) {
      val square1: DoubleLayer = square0 * square0
      if (square1.forward.each.value.data > 0.5) {
        square1 * square1
      } else {
        square1
      }
    } else {
      square0
    }
  } else {
    weight
  }
})

以上动态神经网络中调用了3次forward，但却无处可以调用对应的release。这样写法会导致引用计数没有机会清零，从而导致无法反向传播。

方案四：自动管理引用计数的ResourceT

我们实现了一套ResourceT Monad Transformer，可以用来自动生成对应与forward的release。

在这一方案中，我们把DoubleLayer.forward的类型改为了ResourceT[UnitContinuation, Tape[Double, Double]]。

trait DeepLearning[Differentiable] {
  type Data
  type Delta
  def forward(differentiable: Differentiable): ResourceT[UnitContinuation, Tape[Data, Delta]]
}

ResourceT是个opacity type alias，底层实际类型是UnitContinuation[Resource[Tape[Double, Double]]]。

UnitContinuation是DeepLearning.scala使用的异步编程库。加入UnitContinuation还可以解决forward次数统计的问题。因为UnitContinuation是个惰性计算的Monad，要等到实际运行前才会计算forward次数，这样一来，在动态神经网络中，没有经过的分支就没有开销，运算量小了很多。

由于ResourceT是个Monad，所以我们可以把动态神经网络改成这样：

val weight: DoubleLayer = DoubleWeight(math.random)
def diamondComputationalGraph: DoubleLayer = {
  val resourcet: ResourceT[UnitContinuation, Tape[Double, Double]] = weight.forward
  resourcet.flatMap { tape =>
    if (tape.data > 0.5) {
      val square0: DoubleLayer = weight * weight
      square0.forward.flatMap { tape0 =>
        if (tape0.data > 0.5) {
          val square1: DoubleLayer = square0 * square0
          square1.forward.flatMap { tape1 =>
            if (tape1.data > 0.5) {
              (square1 * square1).forward
            } else {
              square1.forward
            }
          }
        } else {
          square0.forward
        }
      }
    } else {
      weight.forward
    }
  }
}

在神经网络与函数式编程（四）Monadic Deep Learning中，我们介绍过flatMap由Monad.bind实现，作用是生成一个两阶段计算图。

具体的resourcet.flatMap会转发到Monad[ResourceT[UnitContinuation, ?]].bind。

ResourceT作为一个Monad Transformer，其bind实现会额外增加一个功能，把release收集起来：

def bind[A, B](rfa: ResourceT[F, A])(f: (A) => ResourceT[F, B]): ResourceT[F, B] = {
  val ResourceT(fa) = rfa
  fa.flatMap { resourceA =>
    val ResourceT(fb) = f(resourceA.value)
    fb.map { resourceB =>
      val release = resourceB.release >> resourceA.release
      Resource(resourceB.value, release)
    }
  }
  ResourceT(result)
}

这样一来，当用户调用resourcet.flatMap时，返回的ResourceT所包含的release操作由两阶段的release组成：

val release = resourceB.release >> resourceA.release

通过引入ResourceT来管理Resource，我们把release的引用关系自动记录下来了。在每次forward.flatMap时记录引用，最后，在整个网络train时再触发release，所以ResourceT的用户不需要手动调用release，解决了动态神经网络的引用计数问题。

方案五：增加异常处理的Monad Transformer

DeepLearning.scala除了需要引用计数之外，还需要异常处理。和ResourceT类似，这也可以通过增加一层TryT Monad Transformer实现。

这样一来，forward返回的类型变得更加复杂了。

trait DeepLearning[Differentiable] {
  type Data
  type Delta
  def forward(differentiable: Differentiable): TryT[ResourceT[UnitContinuation, ?], Tape[Data, Delta]]
}

方案六：简化Monad类型

在加入支持异步、引用计数管理之后，我们的计算图类型已经非常复杂，变成了TryT[ResourceT[UnitContinuation, ?], Tape[Data, Delta]]，这回给用户带来不必要的学习负担。所以我们创建了一个Do类型。

Do是个opacity type alias。对于任何Do[A]，底层实际类型是TryT[ResourceT[UnitContinuation, ?], A]。

然后我们可以把所有用到TryT、ResourceT之处，改为使用Do：

trait DeepLearning[Differentiable] {
  type Data
  type Delta
  def forward(differentiable: Differentiable): Do[Tape[Data, Delta]
}

方案六是DeepLearning.scala 2.0最终的计算图表示方式。

结论

在DeepLearning.scala 2.0中，我们用Monad表示计算图，把每一个计算图上的特性实现为一个Monad或者一个Monad Transformer，比如：

• 异步计算和并行计算：UnitContinuation
• 反向传播：ResourceT
• 异常处理：TryT

然后用Do把这些Monad Transformer封装起来。最终，DeepLearning.scala的用户并不需要知道Do的内部构造，就能够方便的使用Do构造出动态计算图，自动获得上述所有功能。

在本篇文章中，我们揭示了在Do的简单用法背后隐藏的层层Monad。在下一篇文章中，我将揭示DeepLearning.scala另一个功能，插件系统，背后隐藏的底层机制。