Differentiable Swift のシャープなエッジ

Differentiable Swift は、使いやすさの点で大きな進歩を遂げました。ここでは、まだ少し明らかではない部分について注意してください。進歩が進むにつれて、このガイドはますます小さくなり、特別な構文を必要とせずに微分可能なコードを記述できるようになります。

ループ

ループは微分可能ですが、知っておくべき詳細が 1 つだけあります。ループを記述するときは、ループ対象を指定するビットをwithoutDerivative(at:)でラップします。

var a: [Float] = [1,2,3]

例えば:

for _ in a.indices 
{}

になる

for _ in withoutDerivative(at: a.indices) 
{}

または:

for _ in 0..<a.count 
{}

になる

for _ in 0..<withoutDerivative(at: a.count) 
{}

これが必要なのは、 Array.countメンバーが配列に関する導関数に寄与しないためです。配列内の実際の要素のみが微分に寄与します。

上限として整数を手動で使用するループがある場合は、 withoutDerivative(at:)を使用する必要はありません。

let iterations: Int = 10
for _ in 0..<iterations {} //this is fine as-is.

マップとリデュース

mapreduceは、使い慣れているものとまったく同じように動作する特別な微分可能なバージョンがあります。

a = [1,2,3]
let aPlusOne = a.differentiableMap {$0 + 1}
let aSum = a.differentiableReduce(0, +)
print("aPlusOne", aPlusOne)
print("aSum", aSum)
aPlusOne [2.0, 3.0, 4.0]
aSum 6.0

配列の添字セット

配列添字セット ( array[0] = 0 ) はそのままでは微分できませんが、次の拡張機能を貼り付けることができます。

extension Array where Element: Differentiable {
    @differentiable(where Element: Differentiable)
    mutating func updated(at index: Int, with newValue: Element) {
        self[index] = newValue
    }

    @derivative(of: updated)
    mutating func vjpUpdated(at index: Int, with newValue: Element)
      -> (value: Void, pullback: (inout TangentVector) -> (Element.TangentVector))
    {
        self.updated(at: index, with: newValue)
        return ((), { v in
            let dElement = v[index]
            v.base[index] = .zero
            return dElement
        })
    }
}

回避策の構文は次のようになります。

var b: [Float] = [1,2,3]

これの代わりに:

b[0] = 17

これを書きます:

b.updated(at: 0, with: 17)

それが機能することを確認してみましょう:

func plusOne(array: [Float]) -> Float{
  var array = array
  array.updated(at: 0, with: array[0] + 1)
  return array[0]
}

let plusOneValAndGrad = valueWithGradient(at: [2], in: plusOne)
print(plusOneValAndGrad)
(value: 3.0, gradient: [1.0])

この回避策を使用しないと、 Differentiation of coroutine calls is not yet supportedというエラーが発生します。この回避策を不要にする進捗状況を確認するためのリンクは次のとおりです: https://bugs.swift.org/browse/TF-1277 (配列を実行するときに舞台裏で呼び出される Array.subscript._modify について説明しています)添字セット)。

Float <-> Double変換

FloatDoubleの間で切り替える場合、それらのコンストラクターはまだ微分可能ではありません。これは、 FloatからDoubleに微分的に移行できる関数です。

(以下のコードでFloatDouble切り替えると、 DoubleからFloatに変換する関数が得られます。)

他の実数値型に対しても同様のコンバータを作成できます。

@differentiable
func convertToDouble(_ a: Float) -> Double {
    return Double(a)
}

@derivative(of: convertToDouble)
func convertToDoubleVJP(_ a: Float) -> (value: Double, pullback: (Double) -> Float) {
    func pullback(_ v: Double) -> Float{
        return Float(v)
    }
    return (value: Double(a), pullback: pullback)
}

使用例を次に示します。

@differentiable
func timesTwo(a: Float) -> Double {
  return convertToDouble(a * 2)
}
let input: Float = 3
let valAndGrad = valueWithGradient(at: input, in: timesTwo)
print("grad", valAndGrad.gradient)
print("type of input:", type(of: input))
print("type of output:", type(of: valAndGrad.value))
print("type of gradient:", type(of: valAndGrad.gradient))
grad 2.0
type of input: Float
type of output: Double
type of gradient: Float

超越関数およびその他の関数 (sin、cos、abs、max)

多くの超越関数やその他の一般的な組み込み関数は、 FloatDoubleに対してすでに微分可能になっています。 Doubleの場合はFloatよりも数が少なくなります。どちらにも対応していないものもあります。したがって、必要なものがまだ提供されていない場合に備えて、必要なものを作成する方法のアイデアを提供するために、手動による派生定義をいくつか示します。

pow (派生的な説明についてはリンクを参照)

import Foundation

@usableFromInline
@derivative(of: pow) 
func powVJP(_ base: Double, _ exponent: Double) -> (value: Double, pullback: (Double) -> (Double, Double)) {
    let output: Double = pow(base, exponent)
    func pullback(_ vector: Double) -> (Double, Double) {
        let baseDerivative = vector * (exponent * pow(base, exponent - 1))
        let exponentDerivative = vector * output * log(base)
        return (baseDerivative, exponentDerivative)
    }

    return (value: output, pullback: pullback)
}

最大

@usableFromInline
@derivative(of: max)
func maxVJP<T: Comparable & Differentiable>(_ x: T, _ y: T) -> (value: T, pullback: (T.TangentVector)
  -> (T.TangentVector, T.TangentVector))
{
    func pullback(_ v: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
        if x < y {
            return (.zero, v)
        } else {
            return (v, .zero)
        }
    }
    return (value: max(x, y), pullback: pullback)
}

腹筋

@usableFromInline
@derivative(of: abs)
func absVJP<T: Comparable & SignedNumeric & Differentiable>(_ x: T)
  -> (value: T, pullback: (T.TangentVector) -> T.TangentVector)
{
    func pullback(_ v: T.TangentVector) -> T.TangentVector{
        if x < 0 {
            return .zero - v
        }
        else {
            return v
        }
    }
    return (value: abs(x), pullback: pullback)
}

sqrt (導関数の説明についてはリンクを参照)

@usableFromInline
@derivative(of: sqrt) 
func sqrtVJP(_ x: Double) -> (value: Double, pullback: (Double) -> Double) {
    let output = sqrt(x)
    func pullback(_ v: Double) -> Double {
        return v / (2 * output)
    }
    return (value: output, pullback: pullback)
}

これらが機能することを確認してみましょう。

let powGrad = gradient(at: 2, 2, in: pow)
print("pow gradient: ", powGrad, "which is", powGrad == (4.0, 2.772588722239781) ? "correct" : "incorrect")

let maxGrad = gradient(at: 1, 2, in: max)
print("max gradient: ", maxGrad, "which is", maxGrad == (0.0, 1.0) ? "correct" : "incorrect")

let absGrad = gradient(at: 2, in: abs)
print("abs gradient: ", absGrad, "which is", absGrad == 1.0 ? "correct" : "incorrect")

let sqrtGrad = gradient(at: 4, in: sqrt)
print("sqrt gradient: ", sqrtGrad, "which is", sqrtGrad == 0.25 ? "correct" : "incorrect")
pow gradient:  (4.0, 2.772588722239781) which is correct
max gradient:  (0.0, 1.0) which is correct
abs gradient:  1.0 which is correct
sqrt gradient:  0.25 which is correct

次のようなものが必要であることを警告するコンパイラ エラーは次のとおりです。 Expression is not differentiable. Cannot differentiate functions that have not been marked '@differentiable' and that are defined in other files

KeyPathの添え字付け

KeyPathの添字 (取得または設定) はそのままでは機能しませんが、追加できる拡張機能がいくつかあり、回避策の構文を使用できます。ここにあります:

https://github.com/tensorflow/swift/issues/530#issuecomment-687400701

この回避策は他の回避策よりも少し醜いです。これは、Differentiable および AdditiveArithmetic に準拠する必要があるカスタム オブジェクトに対してのみ機能します。 .tmpメンバーと.read()関数を追加する必要があり、 KeyPathサブスクリプトの取得を実行するときに中間ストレージとして.tmpメンバーを使用します (リンクされたコードに例があります)。 KeyPath添字セットは.write()関数で非常に簡単に動作します。