`getindex` thunk: `add!` disagrees with `val`; `sum(f, xs)` accumulates closure tangents quadratically

オープン
  • 何か問題がありますか？
#863 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
25/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
領域
machine-learning

調査の方向性

Start with the thunked_∇getindex reverse rule for getindex, comparing what its add! does with unthunk(dy) against what val does through ProjectTo(x). Then look at the general branch of rrule for sum(f, xs) and its sum(first, ...) accumulation, which needs an add!! that recurses into Tangent fields. Done means the reproducers in the issue give the same result from add! as from unthunk, and the quadratic allocation case runs in linear time. The design choices for the fix need maintainer input first.

索引モデルが issue の本文から書いたものです。

説明

Two things in the reverse rules keep an AD from accumulating getindex cotangents in place. That makes x[i] in a loop quadratic (FluxML/Zygote.jl#365). Zygote now accumulates in place (FluxML/Zygote.jl#1738), but these two cases stay quadratic because of the rules here. Both are unchanged on master (v1.74.8).

1. thunked_∇getindex: add! adds raw dy, val projects it

val is ∇getindex(x, dy, inds...), which ends in ProjectTo(x)(dx), but add! adds unthunk(dy) as is. So t.add!(dx) can differ from dx + unthunk(t), or throw:

using ChainRules, ChainRulesCore
x = [1.0, 2.0, 3.0]
dx = rrule(getindex, x, 1)[2](1.0 + 2.0im)[2]
unthunk(dx)                    # [1.0, 0.0, 0.0]
dx.add!(zeros(3))              # InexactError
dx.add!(zeros(ComplexF64, 3))  # [1.0 + 2.0im, 0, 0]

dx = rrule(getindex, x, 1:2)[2]([1.0im, 2.0im])[2]
dx.add!(zeros(3))              # InexactError, though unthunk(dx) == zeros(3)

x32 = Float32[1, 2, 3]
dx = rrule(getindex, x32, 2)[2](0.1)[2]
dx.add!(zeros(3))[2] == unthunk(dx)[2]  # false: 0.1 vs Float32(0.1)

So an AD can't call add! without knowing dy, which is hidden in the closure, and has to fall back to unthunk. For a slice that builds a dense array the size of the whole x each time. A loop over the columns of a 100×1000 matrix then allocates 767 MiB in Zygote.

Suggested fix: make add! add the same projected value as val, at least for plain Array{<:Number}, where ProjectTo(x) is elementwise. For example, dx -> ∇getindex!(dx, _project_dy(x, unthunk(dy)), Base.to_indices(x, inds)...), where _project_dy applies ProjectTo(zero(eltype(x))) to dy, elementwise for slices. Structured x could keep dx .+= unthunk(val). With that patch, Zygote calling add! when the inferred eltype of val fits its buffer takes that loop from 767 MiB to 4.4 MiB.

A related cost: in the last example unthunk(dx) is a dense Vector{Float32}, not a OneElement. Projecting the OneElement onto a different eltype materializes it, so mixed-precision scalar indexing is O(length(x)) per access. Projecting dy before building the OneElement would avoid that.

2. sum(f, xs) adds closure tangents with sum(first, df_and_dxs)

In the general branch of rrule(::RuleConfig{>:HasReverseMode}, ::typeof(sum), f, xs), each + of two Tangents whose fields are InplaceableThunks unthunks and copies the whole array. For a closure that captures x and reads x[i], that is O(n²):

using Zygote
x = rand(10_000)
@allocated gradient(x -> sum(i -> x[i]^2, eachindex(x)), x)  # ≈ 766 MiB

Generators hit the same path, since sum(x[i]^2 for i in ...) is lowered to sum(f, iter). Accumulating df into a partial sum the rule owns, instead of sum(first, ...), would make it linear. That needs an add!! that recurses into Tangent fields, since add!! falls back to + for Tangent. Zygote can route around this through its own map adjoint (766 MiB → 3.8 MiB with #1738), but other ADs using this rule would still pay.

🧔‍♂️🤖 Assisted by AI

主要言語
Julia
スター
479
フォーク
100
平均マージ
1日 3時間
マージ済み PR（30日）
13

環境構築

このプロジェクトには開発コンテナ、Dockerfile、コントリビューションガイドがありません。まず README を読み、一般的な手順ははじめてのコントリビューションガイドを参照してください。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

JuliaDiff/ChainRules.jl のほかの issue

JuliaDiff/ChainRules.jl の issue をすべて見る

似ている issue

Julia の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。