Острые края в дифференцируемом Swift

Дифференцируемый Swift прошел долгий путь с точки зрения удобства использования. Вот информация о тех частях, которые все еще немного неочевидны. По мере дальнейшего развития это руководство будет становиться все меньше и меньше, и вы сможете писать дифференцируемый код без необходимости использования специального синтаксиса.

Петли

Циклы дифференцируемы, нужно знать только одну деталь. Когда вы пишете цикл, оберните бит, в котором вы указываете, что вы зацикливаете, 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.

Сопоставить и уменьшить

map и reduce есть специальные дифференцируемые версии, которые работают точно так же, как вы привыкли:

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 преобразование

Если вы переключаетесь между Float и Double , их конструкторы уже не дифференцируемы. Вот функция, которая позволит вам дифференцированно перейти от Float к Double .

(Поменяйте местами Float и Double в приведенном ниже коде, и вы получите функцию, которая преобразует 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)

Многие трансцендентные и другие общие встроенные функции уже стали дифференцируемыми для Float и Double . Их меньше для 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

Этот обходной путь немного уродливее других. Он работает только для пользовательских объектов, которые должны соответствовать дифференцируемым и аддитивным арифметическим. Вам необходимо добавить член .tmp и функцию .read() , и вы используете член .tmp в качестве промежуточного хранилища при получении индекса KeyPath (в связанном коде есть пример). Наборы индексов KeyPath довольно просто работают с функцией .write() .