Автоматическое дифференцирование с использованием Zygote.jl

Flux повторно экспортирует gradient из Zygote и использует эту функцию внутри train! для дифференцирования модели. Для Zygote существует собственная документация, в которой, в частности, перечислены некоторые важные ограничения.

Явный стиль

Предпочтительным способом использования Zygote и единственным способом применения большинства других пакетов AD является явное указание функции и ее аргументов.

# Zygote.gradient — Method

gradient(f, args...)

Возвращает кортеж, содержащий ∂f/∂x для каждого аргумента x, производную (для скалярного x) или градиент. Если градиент не определен, ∂f/∂x будет иметь значение nothing.

f(args...) должно быть вещественным числом, см. описание jacobian для вывода массива.

См. также описание withgradient для сохранения значения f(args...) и pullback для значения и обратного распространения.

julia> gradient(*, 2.0, 3.0, 5.0)
(15.0, 10.0, 6.0)

julia> gradient(x -> sum(abs2,x), [7.0, 11.0, 13.0])
([14.0, 22.0, 26.0],)

julia> gradient([7, 11], 0, 1) do x, y, d
         p = size(x, d)
         sum(x.^p .+ y)
       end
([14.0, 22.0], 2.0, nothing)

# Zygote.withgradient — Method

withgradient(f, args...)
withgradient(f, ::Params)

Возвращает значение функции и gradient как именованный кортеж.

julia> y, ∇ = withgradient(/, 1, 2)
(val = 0.5, grad = (0.5, -0.25))

julia> ∇ == gradient(/, 1, 2)
true

Позволяет захватывать вспомогательные выводы в дополнение к скаляру, используемому градиентом (gradient). Для этого функция f должна возвращать кортеж или именованный кортеж. Затем метод вычисляет +grad = gradient(first∘f, args…) возвращает все значение+val = f(args…)`:

julia> withgradient([1,2,4]) do x
          z = 1 ./ x
          sum(z), z  # здесь z является вспомогательным выводом
       end
(val = (1.75, [1.0, 0.5, 0.25]), grad = ([-1.0, -0.25, -0.0625],))

julia> withgradient(3.0, 4.0) do x, y
          (div = x/y, mul = x*y)
       end
(val = (div = 0.75, mul = 12.0), grad = (0.25, -0.1875))

Также поддерживает неявный режим:

julia> w = [3.0];

julia> res = withgradient(() -> sum(abs2, w), Params([w]))
(val = 9.0, grad = Grads(...))

julia> res.grad[w]
1-element Vector{Float64}:
 6.0

# Zygote.jacobian — Method

jacobian(f, args...) -> Tuple

Для каждого массива a ∈ args возвращает матрицу с Ja[k,i] = ∂y[k]/∂a[i], где y = f(args...) обычно является вектором. Массивы более высоких измерений обрабатываются для вывода как vec(a) или vec(y).

Для скалярного x::Number ∈ args результатом является вектор Jx[k] = ∂y[k]/∂x, тогда как для скалярного y все результаты содержат только одну строку.

При любом другом типе аргумента результат не будет получен, даже если gradient будет работать.

Этот якобиан обратного режима должен вычислять откат один раз для каждого элемента y. Обычно это эффективно только в том случае, если значение length(y) невелико по сравнению с length(a), в противном случае лучше использовать прямой режим.

См. также описание withjacobian, hessian и hessian_reverse.

Примеры

julia> jacobian(a -> 100*a[1:3].^2, 1:7)[1]  # первый индекс (строки) является выводом
3×7 Matrix{Int64}:
 200    0    0  0  0  0  0
   0  400    0  0  0  0  0
   0    0  600  0  0  0  0

julia> jacobian((a,x) -> a.^2 .* x, [1,2,3], 1)  # скалярный аргумент имеет векторный якобиан
([2 0 0; 0 4 0; 0 0 6], [1, 4, 9])

julia> jacobian((a,d) -> prod(a, dims=d), [1 2; 3 4; 5 6], 2)
([2 0 … 0 0; 0 4 … 3 0; 0 0 … 0 5], [0, 0, 0])

Для аргументов любого типа, кроме Number и AbstractArray, результатом является nothing.

julia> jacobian((a,s) -> a.^length(s), [1,2,3], "str")
([3 0 0; 0 12 0; 0 0 27], nothing)

julia> jacobian((a,t) -> sum(a .* t[1]) + t[2], [1,2,3], (4,5))
([4 4 4], nothing)

julia> gradient((a,t) -> sum(a .* t[1]) + t[2], [1,2,3], (4,5))  # градиент понимает кортеж
([4 4 4], (6, 1))

# Zygote.withjacobian — Method

withjacobian(f, args...)

Возвращает значение f(args...) и jacobian как именованный кортеж.

julia> withjacobian(cumsum, [1,2,3])
(val = [1, 3, 6], grad = ([1 0 0; 1 1 0; 1 1 1],))

# Zygote.hessian — Function

hessian(f, x)

Строит гессиан ∂²f/∂x², где x — это вещественное число или массив, а f(x) — вещественное число. Если x — массив, результатом будет матрица H[i,j] = ∂²f/∂x[i]∂x[j], использующая линейное индексирование x[i], даже если аргумент имеет более высокое измерение.

При этом используется прямой, а не обратный режим, ForwardDiff, а не Zygote, вызывающий hessian_dual(f, x). См. hessian_reverse с описанием альтернатив Zygote.

См. также описание diaghessian для вычисления только диагональной части.

Примеры

julia> hessian(x -> x[1]*x[2], randn(2))
2×2 Matrix{Float64}:
 0.0  1.0
 1.0  0.0

julia> hessian(x -> sum(x.^3), [1 2; 3 4])  # использует линейное индексирование x
4×4 Matrix{Int64}:
 6   0   0   0
 0  18   0   0
 0   0  12   0
 0   0   0  24

julia> hessian(sin, pi/2)
-1.0

# Zygote.hessian_reverse — Function

hessian_reverse(f, x)

Может быть эквивалентна hessian(f, x), но реализована с использованием режима «обратный-обратный» для Zygote. (Обычно этот способ намного медленнее и с большей вероятностью обнаружит ошибки.)

# Zygote.diaghessian — Function

diaghessian(f, args...) -> Tuple

Диагональная часть гессиана. Возвращает кортеж, содержащий для каждого аргумента x, h той же формы с h[i] = Hᵢᵢ = ∂²y/∂x[i]∂x[i]. Исходное вычисление y = f(args...) должно выдавать вещественное число y.

Для одного векторного аргумента x это эквивалентно (diag(hessian(f,x)),). Как и в случае с hessian, здесь используется ForwardDiff, а не Zygote.

Для аргументов любого типа, кроме Number и AbstractArray, результатом является nothing.

Примеры

julia> diaghessian(x -> sum(x.^3), [1 2; 3 4])[1]
2×2 Matrix{Int64}:
  6  12
 18  24

julia> Diagonal(vec(ans)) == hessian(x -> sum(x.^3), [1 2; 3 4])  # полный гессиан является диагональным
true

julia> diaghessian((x,y) -> sum(x .* y .* y'), [1 22; 333 4], [0.5, 0.666])  # два аргумента массива
([0.0 0.0; 0.0 0.0], [2.0, 8.0])

julia> diaghessian(atan, 1, 2)  # два скалярных аргумента
(-0.16, 0.16)

julia> hessian(xy -> atan(xy[1], xy[2]), [1, 2])  # полный гессиан не является диагональным
2×2 Matrix{Float64}:
 -0.16  -0.12
 -0.12   0.16

Неявный стиль (Flux версий не выше 0.14)

Раньше во Flux использовалось то, что в Zygote называется «неявными» градиентами, описанными здесь в соответствующей документации. Однако в версии Flux 0.15 поддержка этой возможности будет удалена.

Совместимость: Training

В сине-зеленых блоках в разделе, посвященном обучению, описываются изменения, необходимые для обновления стиля старого кода с неявного на явный.

# Zygote.gradient — Method

gradient(f, args...)

f(args...) должно быть вещественным числом, см. описание jacobian для вывода массива.

См. также описание withgradient для сохранения значения f(args...) и pullback для значения и обратного распространения.

julia> gradient(*, 2.0, 3.0, 5.0)
(15.0, 10.0, 6.0)

julia> gradient(x -> sum(abs2,x), [7.0, 11.0, 13.0])
([14.0, 22.0, 26.0],)

julia> gradient([7, 11], 0, 1) do x, y, d
         p = size(x, d)
         sum(x.^p .+ y)
       end
([14.0, 22.0], 2.0, nothing)

# Zygote.Params — Type

Params([A, B])

Контейнер для неявных параметров, используемый при дифференцировании функции с нулевым аргументом () -> loss(A, B) по отношению к A, B.

# Zygote.Grads — Type

Grads(...)

Подобный словарю контейнер, возвращаемый при взятии градиентов по отношению к неявным параметрам. Для массива W, находящегося в Params([W, A, B...]), градиент имеет значение g[W].

# Zygote.jacobian — Method

jacobian(loss, ::Params)

Как и gradient с неявными параметрами, этот метод принимает функцию с нулевым аргументом и возвращает IdDict-подобный объект, теперь содержащий якобиан для каждого параметра.

Примеры

julia> xs = [1 2; 3 4]; ys = [5,7,9];

julia> Jxy = jacobian(() -> ys[1:2] .+ sum(xs.^2), Params([xs, ys]))
Grads(...)

julia> Jxy[ys]
2×3 Matrix{Int64}:
 1  0  0
 0  1  0

julia> Jxy[xs]
2×4 Matrix{Int64}:
 2  6  4  8
 2  6  4  8

ChainRules

Иногда из автоматического дифференцирования необходимо исключить некоторый код или целую функцию. Это можно сделать с помощью ChainRules:

# ChainRulesCore.ignore_derivatives — Function

ignore_derivatives(f::Function)

Указывает системе AD игнорировать градиенты заключенного в оболочку замыкания. Прямое вычисление (прямой проход) выполняется нормально.

ignore_derivatives() do
    value = rand()
    push!(collection, value)
end

Неправильное использование может привести к появлению неправильных градиентов. Например, следующая функция будет иметь нулевые градиенты по отношению к своему аргументу:

function wrong_grads(x)
    y = ones(3)
    ignore_derivatives() do
        push!(y, x)
    end
    return sum(y)
end

ignore_derivatives(x)

Указывает системе AD игнорировать градиенты аргумента. Можно использовать, чтобы исключить ненужное вычисление градиентов.

ignore_derivatives(x) * w

# ChainRulesCore.@non_differentiable — Macro

@non_differentiable(signature_expression)

Вспомогательный макрос, упрощающий объявление того, что метод не является дифференцируемым. Это сокращение для определения frule и rrule, которые возвращают NoTangent() для всех частных (даже для самой s̄elf-частной функции).

Именованные аргументы включать не следует.

julia> @non_differentiable Base.:(==)(a, b)

julia> _, pullback = rrule(==, 2.0, 3.0);

julia> pullback(1.0)
(NoTangent(), NoTangent(), NoTangent())

В сигнатуру можно ввести ограничения по типу:

julia> @non_differentiable Base.length(xs::Union{Number, Array})

julia> frule((ZeroTangent(), 1), length, [2.0, 3.0])
(2, NoTangent())

Этот вспомогательный макрос охватывает только простые общие случаи. Он не поддерживает предложения where. Для них можно объявить rrule и frule напрямую.

Чтобы вручную задать градиент для одной функции, следует определить метод rrule. О том, как это работает, можно узнать в подробной документации по ChainRules.

# ChainRulesCore.rrule — Function

rrule([::RuleConfig,] f, x...)

Выражая x как кортеж (x₁, x₂, ...), а выходной кортеж f(x...) как Ω, возвращает кортеж:

(Ω, (Ω̄₁, Ω̄₂, ...) -> (s̄elf, x̄₁, x̄₂, ...))

Где второе возвращаемое значение — это правило распространения или откат. Она принимает котангенсы, соответствующие выводам (x̄₁, x̄₂, ...), и s̄elf, внутренние значения самой функции (для замыканий).

Если метод, соответствующий rrule(f, xs...), не был определен, возвращается nothing.

Примеры:

скалярная функция с унарным входом и унарным выводом:

julia> x = rand();

julia> sinx, sin_pullback = rrule(sin, x);

julia> sinx == sin(x)
true

julia> sin_pullback(1) == (NoTangent(), cos(x))
true

скалярная функция с бинарным входом и унарным выводом:

julia> x, y = rand(2);

julia> hypotxy, hypot_pullback = rrule(hypot, x, y);

julia> hypotxy == hypot(x, y)
true

julia> hypot_pullback(1) == (NoTangent(), (x / hypot(x, y)), (y / hypot(x, y)))
true

Необязательный параметр RuleConfig позволяет указывать правила rrule только для систем AD, которые поддерживают заданные возможности. Если он не нужен, его можно опустить, и в качестве резервного варианта будет использоваться rrule без него. Так происходит с большинством правил.

См. также описание frule, @scalar_rule, RuleConfig

# ChainRulesCore.frule — Function

frule([::RuleConfig,] (Δf, Δx...), f, x...)

Выражая вывод f(x...) в виде Ω, возвращает кортеж:

(Ω, ΔΩ)

Второе возвращаемое значение — это тангенс по отношению к выводу.

Если метод, соответствующий frule((Δf, Δx...), f, x...), не был определен, возвращается nothing.

Примеры:

скалярная функция с унарным входом и унарным выводом:

julia> dself = NoTangent();

julia> x = rand()
0.8236475079774124

julia> sinx, Δsinx = frule((dself, 1), sin, x)
(0.7336293678134624, 0.6795498147167869)

julia> sinx == sin(x)
true

julia> Δsinx == cos(x)
true

Скалярная функция с унарным входом и бинарным выводом:

julia> sincosx, Δsincosx = frule((dself, 1), sincos, x);

julia> sincosx == sincos(x)
true

julia> Δsincosx[1] == cos(x)
true

julia> Δsincosx[2] == -sin(x)
true

Обратите внимание, что с технической точки зрения в Julia нет функций множественного вывода, только функции, возвращающие один вывод, который является итерируемым, например Tuple. Таким образом, на самом деле это Tangent:

julia> Δsincosx
Tangent{Tuple{Float64, Float64}}(0.6795498147167869, -0.7336293678134624)

Необязательный параметр RuleConfig позволяет указывать правила frule только для систем AD, которые поддерживают заданные возможности. Если он не нужен, его можно опустить, и в качестве резервного варианта будет использоваться frule без него. Так происходит с большинством правил.

См. также описание rrule, @scalar_rule, RuleConfig

# ChainRulesCore.@scalar_rule — Macro

@scalar_rule(f(x₁, x₂, ...),
             @setup(statement₁, statement₂, ...),
             (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
             (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
             ...)

Вспомогательный макрос, который генерирует простые скалярные прямые и обратные правила, используя предоставленные частные производные. В частности, генерирует соответствующие методы для frule и rrule:

function ChainRulesCore.frule((NoTangent(), Δx₁, Δx₂, ...), ::typeof(f), x₁::Number, x₂::Number, ...)
    Ω = f(x₁, x₂, ...)
    $(statement₁, statement₂, ...)
    return Ω, (
            (∂f₁_∂x₁ * Δx₁ + ∂f₁_∂x₂ * Δx₂ + ...),
            (∂f₂_∂x₁ * Δx₁ + ∂f₂_∂x₂ * Δx₂ + ...),
            ...
        )
end

function ChainRulesCore.rrule(::typeof(f), x₁::Number, x₂::Number, ...)
    Ω = f(x₁, x₂, ...)
    $(statement₁, statement₂, ...)
    return Ω, ((ΔΩ₁, ΔΩ₂, ...)) -> (
            NoTangent(),
            ∂f₁_∂x₁ * ΔΩ₁ + ∂f₂_∂x₁ * ΔΩ₂ + ...),
            ∂f₁_∂x₂ * ΔΩ₁ + ∂f₂_∂x₂ * ΔΩ₂ + ...),
            ...
        )
end

Если в f(x₁, x₂, ...) в рамках вызова @scalar_rule отсутствуют ограничения типа, каждому параметру в результирующем определении frule/rrule задается ограничение типа Number. Ограничения также можно указать явным образом, чтобы переопределить ограничение Number, например f(x₁::Complex, x₂), которое будет ограничивать x₁ типом Complex и x₂ типом Number.

В настоящее время это не поддерживает определение для замыканий или функторов. Таким образом, в обратном режиме первым возвращаемым частным, представляющим производную по самой функции, всегда является NoTangent(). А в прямом режиме первый вход в возвращаемый распространитель всегда игнорируется.

Результат f(x₁, x₂, ...) автоматически привязывается к Ω. Это позволяет легко ссылаться на основной результат (как Ω) в выражениях производной или настройки.

В этом макросе предполагается, что комплексные функции голоморфны. В целом для неголоморфных функций frule и rrule должны быть определены вручную.

Если производная равна единице (например, для функций тождества), true можно использовать как самый общий единичный элемент умножения.

Аргумент @setup может быть опущен, если код настройки не требуется. Другими словами:

@scalar_rule(f(x₁, x₂, ...),
             (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
             (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
             ...)

эквивалентно следующему:

@scalar_rule(f(x₁, x₂, ...),
             @setup(nothing),
             (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
             (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
             ...)

Примеры приведены в каталоге rulesets ChainRules.

См. также описание frule, rrule.

# ChainRulesCore.NoTangent — Type

NoTangent() <: AbstractZero

Этот тангенс указывает, что производная не существует. Это тип тангенса для прямых типов, которые не дифференцируются, например для целых чисел или логических значений (когда они не используются для представления значений с плавающей запятой). Единственный верный способ внести возмущение для таких значений — не изменять их вообще. Как следствие, тип NoTangent функционально идентичен ZeroTangent(), но предоставляет дополнительную семантическую информацию.

Добавление NoTangent() к прямому типу в общем случае неверно: градиентные методы нельзя использовать для оптимизации дискретных переменных. Пакет оптимизации, использующий этот вариант, может захотеть проверить такой случай.

Это не означает, что производная не реализована, просто она не определена математически.

В основном это проявляется в виде производной по аргументам измерения, индекса или размера.

    function rrule(fill, x, len::Int)
        y = fill(x, len)
        fill_pullback(ȳ) = (NoTangent(), @thunk(sum(Ȳ)), NoTangent())
        return y, fill_pullback
    end

# ChainRulesCore.ZeroTangent — Type

ZeroTangent() <: AbstractZero

Аддитивное тождество для тангенсов. По сути, это то же самое, что и 0. Производная ZeroTangent() не распространяется через прямую функцию.