Category Archives: Julia

Progress on Automated Reasoning for Catlab with Metatheory.jl Egraphs

By: Philip Zucker

Re-posted from: https://www.philipzucker.com/metatheory-progress/

Oh cursed equation! How you haunt me!

I’ve been working on mechanizing equational reasoning for categories for a while now with not so much success.

The Egg paper came out and was very compelling to me. It is an implementation with novel twists of egraphs, a useful data structure for equational reasoning. I tried my hand at implementing the egraph data structure here with the purpose of using it for Catlab reasoning in mind.

But Alessandro Cheli has come along and really taken things to the next level with his Metatheory.jl library. The package includes a significantly improved egraph implementation, plus a classical rewriting system. It also has some really interesting ideas in regards to homoiconicity and theory composition. Maybe it’s useful to you. Check it out.

The equational reasoning that occurs in Catlab is compelling because it feels like high school algebra, but describes significantly more powerful systems of reasoning. It also has a beautiful and intuitive diagrammatic interpretation.

Understanding how to take a system like egraph rewriting, which in examples seems at the moment best suited for rewriting ordinary algebraic expressions, and tweak it to work with the more complicated logic of Catlab has been a challenge.

There are a couple problems.

  1. The typing context contains important information. Consider the rule $f \rightarrow id(a) \cdot f $. Looking at only the syntax of this rule, a is invented on the right hand side. This information actually comes from the type of f::hom(a,b)
  2. We need to strip off overloading? $\otimes$ is a symbol used in catlab for both the binary operator on objects and morphisms, as is customary. I guess this might fine? It makes me squeamish.
  3. Some rules require guarding on the type before execution. For most of the rules in Catlab, if the left hand side of the rule is well-typed, then so is the right. This is not the case in general though.
    Consider the trick question: Can you apply the interchange law (f ⊗ g) ⋅ (h ⊗ k) => (f ⋅ h) ⊗ (g ⋅ k) to the term (f ⊗ id(a)) ⋅ (id(b) ⊗ k)?
    No you can’t. For in my example, I have secretly decided that f actually has type b ⊗ c -> b ⊗ c and k has type c ⊗ a -> c ⊗ a.
    The other direction is always fine though. Given (f ⋅ h) ⊗ (g ⋅ k) is well typed, so is (f ⊗ g) ⋅ (h ⊗ k)

I thought way back that the types could be stripped and everything would be fine, like if you strip the types from the simply typed lambda calculus. This is not the case. I think this misconception is due to my inexperience doing hand calculations of this type (which almost no one does. String diagrams rule), plus I think you “code the happy path” when doing it by hand and may not even notice a situation where just shallow pattern matching on the syntax would allow you do do something verboten, since you’re smart enough to not make that dumb mistake.

My initial thinking this time around was to take a type tagging approach https://people.mpi-inf.mpg.de/~jblanche/mono-trans.pdf. This is a standard way of encoding more sophisticated typed logics into simpler typed ones. Type tagging replaces every term ever used with a more complex term that contains the type. In a curious analogy, it is the analog of dynamic typing vs static typing. In dynamically typed languages, all values carry with them tags describing which type they are at runtime.

As an example, the identity id(a) could be tagged as type(id(type(a,ob)), hom(typ(a,ob), type(a,ob))) or as hom(id(ob(a)), ob(a), ob(a)). You can see how verbose this is. It is not ideal both because it is unacceptable to write by hand, and because this extra structure is just more junk to choke your solver. But you can see that now the necessary types to fill in missing pieces of the right hand side is always available.

I tried this through some arcane metaprogramming methods involving Julog and it seemed too slow. Alessandro and James Fairbanks have a more elegant looking version of this here.

It feels insane that you need all this overhead when it is not necessary when working by hand. Many of the rules do not have this rhs invention problem or the typed guarding requirement. For cartesian categories, only about 5 of 30 rules have this condition. $(f \cdot g) \cdot h \rightarrow (f \cdot g) \cdot h$ is totally fine for example.

So here is idea 1, replace any new derived information on the right hand side with “function calls” that find that information. For our example, this becomes f => id(dom(type(f))) ⋅ f. Then we need to internalize the typing rules into the rewrite system

typing_rules = @theory begin
    dom(hom(a,b)) => a
    cod(hom(a,b)) => b
    type(id(a)) => hom(a,a)
    type(f ⋅ g) => hom( dom(type(f)) , cod(type(g)) )
    type(f ⊗ₘ g) => hom( dom(type(f)) ⊗ₒ dom(type(g)) , cod(type(f)) ⊗ₒ cod(type(g)))
    type(a ⊗ₒ b) => :ob
    type(munit()) => :ob
    type(σ(a,b)) => hom(a ⊗ₒ b, b ⊗ₒ a)
    type(⋄(a)) => hom(a, munit())
    type(Δ(a)) => hom(a, a ⊗ₒ a)
    type(pair(f,g)) => hom(dom(type(f)), cod(type(f)) ⊗ₒ cod(type(g)))
    type(proj1(a,b)) => hom(a ⊗ₒ b, a)
    type(proj2(a,b)) => hom(a ⊗ₒ b, b)
end

Idea two is to use Metatheory “guards” to enforce type correctness of the left hand side before a rule can fire. As part of this, the necessary type queries are added to the egraph. Alessandro had this idea of just returning the original expression to enforce a guard, which I think is pretty clever. Here is how I think the guard works for the interchange law. This will be cleaner with an upcoming Metatheory macro.

    (f ⊗ₘ p) ⋅ (g ⊗ₘ q) |>
    # future metatheory macros will clean this up
    begin
        fcod = addexpr!(_egraph, :(cod(type($f))))
        gdom = addexpr!(_egraph, :(dom(type($g))))
        pcod = addexpr!(_egraph, :(cod(type($p))))
        qdom = addexpr!(_egraph, :(dom(type($q))))
        if (fcod == gdom && pcod == qdom)
            :(($f ⋅ $g) ⊗ₘ ($p ⋅ $q))
        else
            :(($f ⊗ₘ $p) ⋅ ($g ⊗ₘ $q))
        end
    end

You can see that even making the query about whether the types are right adds them to the egraph. It may be currently unknown what the types are, but by adding them, eventually the typing rules will fire and this rule will have its chance again. This also isn’t as bad as it may appear performance wise, because the egraph in a sense memoizes these typing calls. See the Fibonacci example by Mason Protter which has non exponential complexity despite the naive implementation https://github.com/0x0f0f0f/Metatheory.jl/blob/master/test/test_fibonacci.jl

It is my belief that both of these points (internalization of the typing rules and guards) can be mechanized to produce a sound rule system for an arbitrary GAT. I could easily be wrong. The rough picture I suggest is that you need to infer by looking at the lhs of each rule what typing constraints are implied. If you do not imply the entire typing context, you need to add as many extra guards to ensure it as needed. For example, the well typed-ness of (f ⊗ₘ p) ⋅ (g ⊗ₘ q) implies $cod(f) \otimes cod(p) = dom(g) \otimes dom(q)$ but this does not imply the condition $cod(f) = dom(g) \land cod(p) = dom(q)$ so they must be added.

EGraphs are confusing to reason about. But 2 points:

  1. I think it is important that the Egraph inductively contains only well typed terms in some sense. This should be definable and checkable in a precise manner.
  2. I think it is important to have a semantic model in mind when working with egraphs. Pure syntax considerations lead you astray. EGraphs are more than just a bag of terms.

Here is a full script to run the derivation pair(proj1(a,b),proj2(a,b)) = id(a ⊗ₒ b). I can’t get it to go end to end without assistance yet. We’ll get there. I believe. The real holy grail is not proofs, but simplifications. This also works, but not yet at the depth I’m hoping for.

Evan Patterson’s explanation of the proof has been absolutely vital https://julialang.zulipchat.com/#narrow/stream/230248-catlab.2Ejl/topic/Automatic.20Theorem.20Proving/near/211396362. Zulip is nice for this.

using Pkg
Pkg.activate("metacat")
Pkg.add(url = "https://github.com/0x0f0f0f/Metatheory.jl.git")
using Metatheory
using Metatheory.EGraphs
using Metatheory.Util

# See this link for the catlab axioms
# https://github.com/AlgebraicJulia/Catlab.jl/blob/ce2fde9c63a8aab65cf2a7697f43cd24e5e00b3a/src/theories/Monoidal.jl#L127

cat_rules = @theory begin
    f  id(b) => f
    id(a)  f => f

    #f => f ⋅ id(dom(type(f)))
    #f => id(cod(type(f))) ⋅ f

    a ⊗ₒ munit() => a
    munit() ⊗ₒ a => a

    #a => a ⊗ₒ munit() 
    #a => munit() ⊗ₒ a

    f  (g  h) == (f  g)  h


end

monoidal_rules = @theory begin
    id(munit()) ⊗ₘ f => f
    f ⊗ₘ id(munit()) => f
    a ⊗ₒ (b ⊗ₒ c) == (a ⊗ₒ b) ⊗ₒ c
    f ⊗ₘ (h ⊗ₘ j) == (f ⊗ₘ h) ⊗ₘ j
    id(a ⊗ₒ b) == id(a) ⊗ₘ id(b)

    (f ⊗ₘ p)  (g ⊗ₘ q) |>
    # future metatheory macros will clean this up
    begin
        fcod = addexpr!(_egraph, :(cod(type($f))))
        gdom = addexpr!(_egraph, :(dom(type($g))))
        pcod = addexpr!(_egraph, :(cod(type($p))))
        qdom = addexpr!(_egraph, :(dom(type($q))))
        if (fcod == gdom && pcod == qdom)
            :(($f  $g) ⊗ₘ ($p  $q))
        else
            :(($f ⊗ₘ $p)  ($g ⊗ₘ $q))
        end
    end

    (f  g) ⊗ₘ (p  q) => (f ⊗ₘ p)  (g ⊗ₘ q)
end



sym_rules = @theory begin
    σ(a, b)  σ(b, a) == id(a ⊗ₒ b)
    (σ(a, b) ⊗ₘ id(c))  (id(b) ⊗ₘ σ(a, c)) == σ(a, b ⊗ₒ c)
    (id(a) ⊗ₘ σ(b, c))  (σ(a, c) ⊗ₘ id(b)) == σ(a ⊗ₒ b, c)

    (f ⊗ₘ h)  σ(a, b) |> begin
        fcod = addexpr!(_egraph, :(cod(type($f)))).id
        hcod = addexpr!(_egraph, :(cod(type($h)))).id
        if fcod == a && hcod == b
            :(σ(dom(type($f)), dom(type($h)))  ($h ⊗ₘ $f))
        else
            :(($f ⊗ₘ $h)  σ($a, $b))
        end
    end


    σ(c, d)  (h ⊗ₘ f) |> begin
        fdom = addexpr!(_egraph, :(dom(type($f)))).id
        hdom = addexpr!(_egraph, :(dom(type($h)))).id
        if fdom == c && hdom == d
            :(($f ⊗ₘ $h)  σ(cod(type($f)), cod(type($h))))
        else
            :(σ($c, $d)  ($h ⊗ₘ $f))
        end
    end

    # these rules arer not catlab
    σ(a, munit()) == id(a)
    σ(munit(), a) == id(a)
    σ(munit(), munit()) => id(munit() ⊗ₒ munit())

end



diag_rules = @theory begin

    Δ(a)  ((a) ⊗ₘ id(a)) == id(a)
    Δ(a)  (id(a) ⊗ₘ (a)) == id(a)
    Δ(a)  σ(a, a) == Δ(a)

    (Δ(a) ⊗ₘ Δ(b))  (id(a) ⊗ₘ σ(a, b) ⊗ₘ id(b)) == Δ(a ⊗ₒ b)


    Δ(a)  (Δ(a) ⊗ₘ id(a)) == Δ(a)  (id(a) ⊗ₘ Δ(a))
    (a ⊗ₒ b) == (a) ⊗ₘ (b)

    Δ(munit()) == id(munit())
    (munit()) == id(munit())
end


cart_rules = @theory begin

    Δ(a)  (f ⊗ₘ k) |> begin
        a_id = find(_egraph, a)
        if (
            addexpr!(_egraph, :(dom(type($f)))).id == a_id &&
            addexpr!(_egraph, :(dom(type($k)))).id == a_id
        )
            :(pair($f, $k))
        else
            :(Δ($a)  ($f ⊗ₘ $k))
        end
    end


    pair(f, k) => Δ(dom(type(f)))  (f ⊗ₘ k)
    proj1(a, b) == id(a) ⊗ₘ (b)
    proj2(a, b) == (a) ⊗ₘ id(b)
    f  (b) => (dom(type(f)))
    # Has to invent f. Hard to fix.
    # ⋄(b) => f ⋅ ⋄(b)  

    f  Δ(b) => Δ(dom(type(f)))  (f ⊗ₘ f)
    Δ(a)  (f ⊗ₘ f) => f  Δ(cod(type(f)))
end





typing_rules = @theory begin
    dom(hom(a, b)) => a
    cod(hom(a, b)) => b
    type(id(a)) => hom(a, a)
    type(f  g) => hom(dom(type(f)), cod(type(g)))
    type(f ⊗ₘ g) => hom(dom(type(f)) ⊗ₒ dom(type(g)), cod(type(f)) ⊗ₒ cod(type(g)))
    type(a ⊗ₒ b) => :ob
    type(munit()) => :ob
    type(σ(a, b)) => hom(a ⊗ₒ b, b ⊗ₒ a)
    type((a)) => hom(a, munit())
    type(Δ(a)) => hom(a, a ⊗ₒ a)
    type(pair(f, g)) => hom(dom(type(f)), cod(type(f)) ⊗ₒ cod(type(g)))
    type(proj1(a, b)) => hom(a ⊗ₒ b, a)
    type(proj2(a, b)) => hom(a ⊗ₒ b, b)
end


rules = typing_rules  cat_rules  monoidal_rules  sym_rules  diag_rules  cart_rules  typing_rules


# A goofy little helper macro
# Taking inspiration from Lean/Dafny/Agda
macro calc(e...)
    theory = eval(e[1])
    e = rmlines(e[2])
    @assert e.head == :block
    for (a, b) in zip(e.args[1:end-1], e.args[2:end])
        println(a, " =? ", b)
        @time println(areequal(theory, a, b; timeout = 40))
    end
end

# Get the Julia motor hummin'
@calc rules begin

    (((a) ⊗ₘ (b))  σ(munit(), munit()))
    (σ(a, b)  ((b) ⊗ₘ (a)))

end 


@calc rules begin
    id(a ⊗ₒ b)
    id(a) ⊗ₘ id(b)
    (Δ(a)  (id(a) ⊗ₘ (a))) ⊗ₘ id(b)
    (Δ(a)  (id(a) ⊗ₘ (a))) ⊗ₘ (Δ(b)  ((b) ⊗ₘ id(b)))
    (Δ(a) ⊗ₘ Δ(b))  ((id(a) ⊗ₘ (a)) ⊗ₘ ((b) ⊗ₘ id(b)))
    (Δ(a) ⊗ₘ Δ(b))  (id(a) ⊗ₘ ((a) ⊗ₘ (b)) ⊗ₘ id(b))
    (Δ(a) ⊗ₘ Δ(b))  (id(a) ⊗ₘ (((a) ⊗ₘ (b))  σ(munit(), munit())) ⊗ₘ id(b))
    (Δ(a) ⊗ₘ Δ(b))  ((id(a) ⊗ₘ (σ(a, b)  ((b) ⊗ₘ (a))) ⊗ₘ id(b)))
    (Δ(a) ⊗ₘ Δ(b))  ((id(a) ⊗ₘ (σ(a, b)  ((b) ⊗ₘ (a))) ⊗ₘ id(b)))
    (Δ(a) ⊗ₘ Δ(b))  ((id(a)  id(a)) ⊗ₘ (σ(a, b)  ((b) ⊗ₘ (a))) ⊗ₘ id(b))
    (Δ(a) ⊗ₘ Δ(b))  ((id(a) ⊗ₘ σ(a, b) ⊗ₘ id(b))  (id(a) ⊗ₘ ((b) ⊗ₘ (a)) ⊗ₘ id(b)))
    Δ(a ⊗ₒ b)  (id(a) ⊗ₘ ((b) ⊗ₘ (a)) ⊗ₘ id(b))
    Δ(a ⊗ₒ b)  (id(a) ⊗ₘ ((b) ⊗ₘ (a)) ⊗ₘ id(b))
    Δ(a ⊗ₒ b)  (proj1(a, b) ⊗ₘ proj2(a, b))
    pair(proj1(a, b), proj2(a, b))
end

# shorter proof also accepted
@calc rules begin
    id(a ⊗ₒ b)
    (Δ(a) ⊗ₘ Δ(b))  ((id(a) ⊗ₘ (σ(a, b)  ((b) ⊗ₘ (a))) ⊗ₘ id(b)))
    pair(proj1(a, b), proj2(a, b))
end

Where to go from here

  • We could just say this is already nice and push forward.
  • You may have noticed I have a few rules commented out. These may or may not present extra difficulties. The unit production rules are probably fine with another kind of guard. The deletion rule in the other direction will never be fine. This rule truly invents f in an essential way.
  • This is all totally unintegrated with catlab itself. A WIP. There are different directions one could go on this. Alessandro has a branch of Metatheory in the work for custom user defined data structures, so it is possible to use Catlab’s native GATExpr maybe. He’s looking for help too btw!
  • I’m kind of intrigued at trying to make a pure egg version of the above, especially so I can compile to wasm
    -String diagrams https://julialang.zulipchat.com/#narrow/stream/230248-catlab.2Ejl/topic/Using.20Metatheory.20for.20String.20Diagrams There is a compelling argument that string diagrams are a preferred representation, normalizing the associativity,commutative, and unit properties of a categorical construction. The useful canonicity of string diagrams is a viewpoint Evan Patterson has put forward. It has been suggested by Alessandro and others that generalizing the EGraph data structure in some sense may be to go. Extensions of the egraph that work natively on something more string diagram-like?
  • More wildly speculative. Co-end calculus? How to deal with binding forms? Daniele Palombi has brought up the coend calculus https://arxiv.org/pdf/1501.02503.pdf This seems like a very interesting application. I think the direct attack on this problem using Metatheory requires understanding how to deal with alpha equivalence in the EGraph Conversation here: https://julialang.zulipchat.com/#narrow/stream/277860-metatheory.2Ejl/topic/Formal.20Proofs.20for.20CT.3F. I personally would like to start attacking the problem from the angle of regular integrals and sums, which are syntactically analogous.
  • Proof Recording and output. The actual proof is there in the egraph rewriting process if we record it. Alessandro sounded interested in this direction and my understanding is it is also in the works for egg
  • Cody has also pointed out to me that he had a project with similar goals but different methods many moons ago https://github.com/codyroux/catnapp. I don’t think this got that far.

What is new in PooledArrays.jl?

By: Blog by Bogumił Kamiński

Re-posted from: https://bkamins.github.io/julialang/2021/03/05/pooledarrays.html

Introduction

Recently PooledArrays.jl 1.2.1 has been released. The most significant change
since 1.0 release is an improvement of performance of basic operations:
getindex, copy, copy!, and copyto!. The effect of the change is
especially significant for PooledArrays that have large pools. This change, is
one of the steps towards making Julia run fast in Database-like ops
benchmark
for joins.

Let me start with the examples and then I will comment on the internals.
The post was tested under Julia 1.6.0-rc1.

The Benchmarks

I will just share a recordng of a Julia session doing the benchmarks. We start
with PooledArrays.jl 1.0.0:

(@v1.6) pkg> activate .
  Activating new environment at `~/Project.toml`

(bkamins) pkg> add [email protected]
   Resolving package versions...
    Updating `~/Project.toml`
  [2dfb63ee] + PooledArrays v1.0.0
    Updating `~/Manifest.toml`
  [9a962f9c] + DataAPI v1.6.0
  [2dfb63ee] + PooledArrays v1.0.0
  Progress [========================================>]  1/1
1 dependency successfully precompiled in 2 seconds (1 already precompiled)

julia> using BenchmarkTools, PooledArrays

julia> x = PooledArray(string.(1:10^6));

julia> @benchmark copy($x)
BenchmarkTools.Trial:
  memory estimate:  37.44 MiB
  allocs estimate:  13
  --------------
  minimum time:     16.109 ms (0.00% GC)
  median time:      17.820 ms (0.00% GC)
  mean time:        29.851 ms (40.25% GC)
  maximum time:     175.307 ms (88.17% GC)
  --------------
  samples:          168
  evals/sample:     1

julia> @benchmark $x[1:1]
BenchmarkTools.Trial:
  memory estimate:  33.63 MiB
  allocs estimate:  12
  --------------
  minimum time:     15.334 ms (0.00% GC)
  median time:      16.929 ms (0.00% GC)
  mean time:        30.654 ms (44.06% GC)
  maximum time:     188.496 ms (89.90% GC)
  --------------
  samples:          164
  evals/sample:     1

In order to assess if the timings are good or bad let us do the same operations
using a plain Vector{String}:

julia> x = string.(1:10^6);

julia> @benchmark copy($x)
BenchmarkTools.Trial:
  memory estimate:  7.63 MiB
  allocs estimate:  2
  --------------
  minimum time:     395.687 μs (0.00% GC)
  median time:      441.842 μs (0.00% GC)
  mean time:        757.598 μs (38.52% GC)
  maximum time:     6.194 ms (90.51% GC)
  --------------
  samples:          6598
  evals/sample:     1

julia> @benchmark $x[1:1]
BenchmarkTools.Trial:
  memory estimate:  96 bytes
  allocs estimate:  1
  --------------
  minimum time:     25.310 ns (0.00% GC)
  median time:      26.013 ns (0.00% GC)
  mean time:        34.384 ns (8.88% GC)
  maximum time:     2.433 μs (98.05% GC)
  --------------
  samples:          10000
  evals/sample:     992

As you can see things are really bad in PooledArrays.jl. Now start a fresh Julia
session and install 1.2.1 release of the package.

(bkamins) pkg> activate .
  Activating environment at `~/Project.toml`

(bkamins) pkg> add [email protected]
   Resolving package versions...
    Updating `~/Project.toml`
  [2dfb63ee] ↑ PooledArrays v1.0.0 ⇒ v1.2.1
    Updating `~/Manifest.toml`
  [2dfb63ee] ↑ PooledArrays v1.0.0 ⇒ v1.2.1
  [9fa8497b] + Future
  [9a3f8284] + Random
  [9e88b42a] + Serialization

julia> using BenchmarkTools, PooledArrays

julia> x = PooledArray(string.(1:10^6));

julia> @benchmark copy($x)
BenchmarkTools.Trial:
  memory estimate:  3.81 MiB
  allocs estimate:  4
  --------------
  minimum time:     893.391 μs (0.00% GC)
  median time:      948.740 μs (0.00% GC)
  mean time:        1.009 ms (5.08% GC)
  maximum time:     103.288 ms (98.55% GC)
  --------------
  samples:          4953
  evals/sample:     1

julia> @benchmark $x[1:1]
BenchmarkTools.Trial:
  memory estimate:  160 bytes
  allocs estimate:  3
  --------------
  minimum time:     47.806 ns (0.00% GC)
  median time:      53.053 ns (0.00% GC)
  mean time:        184.150 ns (44.04% GC)
  maximum time:     55.779 μs (68.58% GC)
  --------------
  samples:          10000
  evals/sample:     976

This looks better. Still you have to pay some cost over a Vector{Sting} but it
is much smaller (the cost is due to the fact that PooledArray constructor
performs some consistency checks of passed data to ensure extra safety).

You can expect that other operations that take a PooledArray or its view and
produce a PooledArray (like or copyto!) to experience similar
speedups.

So, what has changed between 1.0 and 1.2 release of PooledArrays.jl?

The Internals

In order to understand why the speedups were possible one needs to understand
first why the original code was slow. The reason is that PooledArray struct
contains three key fields:

  • refs: a vector of integer references (levels) of a PooledArray;
  • pool: a vector giving a mapping from references to actual values;
  • invpool: a dictionary providing a reverse mapping – from values to references.

As you can see this data structure is quite heavy if the size of pool relative
to the size of the PooledArray is large. In the example above I have shown an
extreme case where they were equal. But even if pool has e.g. 10% of size of
the PooledArray the cost is noticeable.

In PooledArrays.jl 1.0 all thee fields refs, pool and invpool were always
copied when a new PooledArray was created. This was expensive. What
PooeldArrays.jl 1.2 introduces is a well known from R and often asked about in
Julia copy on write behavior. What we now do is that we only copy refs. The
pool and invpool fields are shared across PooledArrays as this is a safe
thing to do as long as you are not changing the set of levels in the pool.

So where does the copy on write happen? PooledArrays.jl is now aware if you
share pool and invpool across several arrays and if this is the case and you
add levels to PooledArray then pool and invpool get copied. So essentially
we have a lazy mechanism that copies them only if needed.

One could ask why do we copy pool and invpool at all. One could just keep
sharing them without having to pay the cost of copying at all. The decision was
guided here by two considerations:

  • in practice new levels are added to PooledArray quite rarely (you mostly do
    it when constructing an initial source PooledArray);
  • it is safer to copy the pool and invpool in consideration of potential
    multi-threaded usages of PooledArray (where many tricky corner cases can
    happen).

Conclusions

The main take away is that you can expect your code using PooledArrays.jl to be
much faster since 1.2 release.

Before I finish let me comment on one important feature of PooledArray object
to keep in mind in the multi-threaded applications I have mentioned above.

It is currently not thread safe to add new levels to the pool
of PooledArray. So the rule is: if you add levels to PooledArray make sure
you are not performing any other operations on it in other threads.

However, you can safely perform any operations in multi-threaded context that do
not change the pool. So, in particular (barring standard considerations of
correct multi-threaded code), you are allowed to use setindex! on
PooledArray as long as you do not add new levels.

Machine-learned preconditioners – Part 1

By: Random blog posts about machine learning in Julia

Re-posted from: https://rkube.github.io/jekyll/update/2021/03/02/neural-networks-and-iterative-solvers.html

What is a preconditioner?

A common task in scientific computing is to solve a system of linear equations

\[Ax = b\]

with a matrix \(A \in \mathcal{R}^{d \times d}\) that gives the coefficients of the
system, a right-hand-side vector \(b \in \mathcal{R}^{d}\),
and a solution vector \(x \in \mathcal{R}^{d}\).

Instead of solving the linear problem directly, one often solves
the preconditioned problem. We write this problem by inserting a one in the
linear system and re-parenthesing:

\(\left( A P^{-1} \right) \left( P x \right) = b\)
In this equation we have changed the coefficient matrix from \(A\) to
\(A P^{-1}\) and the solution vector is now \(P x\) instead of \(x\). To
retrieve the original vector, simply calculate \(x = P^{-1}y\). The matrix
\(P^{-1}\) is called the preconditioner.

The dependence of \(P^{-1}\) on \(x\) is by choice. There is no deeper reason
for why the preconditioner should depend on the initial guess for the iteraion.
We choose to include this dependency here to foreshadow later applications, where
such a dependency may be useful. In practice, this choice here is rather limiting.
Since the matrix \(A_0 P^{-1}\) needs to be positive definite, and by construction
only \(\widetilde{b}\) is positive definite, the method as constructed here is only
valid for \(x \approx 0\).

The Gauss-Seidel method is an iterative method to find the solution of
a linear system. Given an initial guess \(x_0\), update the entries of this
vector following an iterative scheme and after some, or many, iterations, \(x_0\)
solves the equation \(A x_0 = b\).
Wikipedia has a nice
and instructive page on this scheme.

A good preconditioner gives you a converged solution after fewer iterations.
In other words, with a good preconditioner you are closer to the true solution
of the linear system after N iterations than you are with a bad preconditioner.
Choosing a good preconditioner depends on the problem at hand and can become
a dark art. Let’s make it even more dark and train a neural network to be
a preconditioner.

Modeling preconditioner as a neural network

As a first step, let’s model \(P^{-1}\) as a multi-layer perceptron
(technically a single-layer):

\[P^{-1} = \sigma \left( W x + \widetilde{b} \right)\]

with a weight matrix \(W \in \mathcal{R}^{d^2 \times d}\), a bias
\(\widetilde{b} \in \mathcal{R}^{d^2}\), and a ReLU \(\sigma\). With
\(x \in \mathcal{R}^{d}\) the matrix dimensions are chosen such that \(P^{-1}\)
is of dimension \(d^2\). Simply reshape it to \(d \times d\) to have it act like
a matrix.

Here I’m discussing a simple test case and am working with the Gauss-Seidel
iterative solver. For real problem one would probably use a different iterative
algorithm, but it serves as a proof-of-concept. Anyway, for Gauss-Seidel to work,
the system coefficient matrix \(AP^{-1}\) needs to be positive-definite.
How do we do this for our neural network? A simple hack is to consider only
vectors whose entries are close to zero.

That way we get away with requiring only \(\widetilde{b}\) to be
positive-definite. We get the last property by letting
\(\widetilde{b} = b_0 b_0^{T}\) and sampling the entries of \(b_0\) as
\(b_{ij} \sim \mathcal{N}(0, 1)\). In a similar fashion, we sample the entries
of the weight-matrix \(W\) as \(w_{i,j} \sim \mathcal{N}(0,1)\).

How can we train a preconditioner

To make this preconditioner useful, the weights \(W\) and bias term \(b\) need
to be optimized such that the residual after \(N\) iterations is as small
as possible. And this needs to be true for all vectors from a training set.
Using automatic differentiation, we can train \(W\) and \(b\) using gradient
descent like this:

  • Pick an initial guess \(x_0\) with \(x_{0,i} \sim \mathcal{N}(0, 0.01)\)
  • Build the preconditioner \(P^{-1} = \sigma(W, x_0, \widetilde{b})\)
  • Calculate 5 Gauss-Seidel steps. Let’s call the solution here \(y_5\).
  • Un-apply the preconditioner: \(x_5 = P^{-1} y_5\)
  • Calculate the distance to the true solution \(\mathcal{L} = \frac{1}{N} \sum_{i=1}^{N} \left( x_{\text{true}, i} – x_{5,i} \right)^2\)
  • Calculate the gradients \(\nabla_{W} \mathcal{L}\) and
    \(\nabla_{B} \mathcal{L}\)
  • Update the weight matrix and bias using gradient descent: \(W \leftarrow W – \alpha \nabla_{W} \mathcal{L}\) and \(b \leftarrow b – \alpha \nabla_{b} \mathcal{L}\). Here \(\alpha\) is the learning rate

The approach here is to directly back-propagate from the loss function \(\mathcal{L}\), through
the numerical solver, to the parameters of the preconditioner \(W\) and \(b\).
We are not working with an offline training and test-data set, but the data is
taken directly from the numerical calculations. This way we directly capture the
reaction of the numerical solver to updates proposed by gradient descent.

The derivatives \(\nabla_\theta \mathcal{L}\) can be calculate using automatic differentiation
packges, such as Zygote.

Implementation


# Test differentiation through control flow

# Use a iterative conjugate solver with preconditioner

using Random
using Zygote
using LinearAlgebra
using Distributions
using NNlib

I copy-and-pasted the Gauss-Seidel code from wikipedia:

function gauss_seidel(A, b, x, niter)
    # This is from https://en.wikipedia.org/wiki/Gauss%E2%80%93Seidel_method
    x_int = Zygote.Buffer(x)
    x_int[:] = x[:]
    for n  1:niter
        for j  1:size(A, 1)
            x_int[j] = (b[j] - A[j, :]' * x_int[:] + A[j, j] * x_int[j]) / A[j, j]
        end
    end
    return copy(x_int)
end

The block below sets things up.


Random.seed!(1)
dim = 4

# Define a matrix
A0 = [10.0 -1.0 2.0 0.0; -1 11 -1 3; 2 -1 10 -1; 0 3 -1 8]
# Define the RHS
b0 = [6.0; 25.0; -11.0; 15.0]
# This is the true solution
x_true = [1.0, 2.0, -1.0, 1.0]

# Define the size of the training and test set
N_train = 100
N_test = 10

# Define an initial state. Draw this from a narrow distribution around zero
# We need to do this so that the preconditioner eigenvalues are positive
x0 = rand(Normal(0.0, 0.01), (N_train, dim))

# Define an MLP. This will later be our preconditioner
# The output should be a matrix and we work with 2dim as the size for the MLP
W = rand(dim*dim, dim)
# For Gauss-Seidel to work, the matrix A*P⁻¹ needs to be positive semi-definite.
bmat = rand(dim, dim)
# We know that any matrix A*A' is positive semi-definite
bvec = reshape(bmat * bmat', (dim * dim))
# Now Wx + bmat is positive semi-definite if x is very small
P(x, W, b) = NNlib.relu.(reshape(W*x .+ b, (dim, dim)))
# Positive-definite means positive Eigenvalues. We should check this.
@show eigvals(A0*P(x0[0,:], W, bvec))

This function will serve as our loss function. It evaluates the NN preconditioner and then runs
some Gauss-Seidel iterations. Finally, the iterative solution
approximation is transformed back by applying \(P^{-1}\).

function loss_fun(W, bmat, A0, b0, y0, niter=5)
    # W - Weight matrix for NN-preconditioner
    # bmat - Bias vector for NN preconditioner
    # A0: Linear system coefficient matrix
    # b0: RHS of linear system
    # y0 - initial guess for Linear system. Strictly, this is x0. But we call it the same
    # assuming that P⁻¹x0 = x0.
    # niter - Number of Gauss-Seidel iterations to perform
    
    loss = 0.0
    nsamples = size(y0)[1]
    
    for idx ∈ 1:nsamples:
        # Evaluate the preconditioner
        P⁻¹ = P(y0[idx, :], W, reshape(bmat * bmat', (dim * dim)))
        # Initial guess
        # Now we solve A(Px)⁻¹y = rhs for y with 3 Gauss-Seidel iterations
        y_sol = gauss_seidel(A0 * P⁻¹, b0, y0[idx, :], niter)
        # And reconstruct x
        x = P⁻¹ * y_sol
        loss += norm(x - x_true) / length(x)
    end

    return loss / nsamples
end

Now comes the fun part. Zygote calculates the partial derivatives of the
loss function with respect to its input, in our case \(W\) and \(b\).
Given the gradients, we can actually update \(W\) and \(b\).

# Number of epochs to train
num_epochs = 10
loss_arr = zeros(num_epochs)
# Store the Weight and bias matrix at each iteration
W_arr = zeros((size(W)..., num_epochs))
bmat_arr = zeros((size(bmat)..., num_epochs))
# Learning rate
α = 0.005

for epoch ∈ 1:num_epochs
    loss_arr[epoch], grad = Zygote.pullback(loss_fun, W, bmat, A0, b0, x0)
    res = grad(1.0)
    # Gradient descent
    global W -= α * res[1]
    global bmat -= α * res[2]
    # Store W and b
    W_arr[:, :, epoch] = W
    bmat_arr[:, :, epoch] = bmat
end

Finally, let’s evaluate the performance

# This functions returns a vector with the residual of the iterative
# solution to the true solution at each step
function eval_performance(W, bmat, A0, y0, b0, niter=20)
    # Instantiate the preconditioner with the initial guess
    P⁻¹ = P(y0, W, reshape(bmat * bmat', (dim * dim)))
    y_sol = copy(y0)
    loss_vec = zeros(niter)
    for n ∈ 1:niter      
        # Initial guess
        # Now we solve (Px)⁻¹y = rhs for y with 3 Gauss-Seidel iterations
        y_sol = gauss_seidel(A0 * P⁻¹, b0, y_sol, niter)
        # And reconstruct x
        x = P⁻¹ * y_sol
        loss_vec[n] = norm(x - x_true) / length(x)
    end

    return loss_vec
end

# Get the residual at each Gauss-Seidel iteration
sol_err = zeros(20, num_epochs)
for i ∈ 1:num_epochs
    # Calculate the loss at each iteration, averaged over the training data
    sol_avg = zeros(20)
    for idx ∈ 1:N_test
        sol_here[:] += eval_performance(W_arr[:, :, i], bmat_arr[:, :, i], A0, x0, b0)
    end
    sol_err[:, i] = sol_here[:] / N_test
end

Results

A plot says many words, so here we go

Trained preconditioner

The plot shows the average residual to the true solution vector as a function of
Gauss-Seidel iterations. Training for 1 epoch, the residual decreases as a power law
for all 20 GS iterations. The longer we train the preconditioner, the faster
the residual decreases. Remember that we trained only for 5 iterations. But in the plot
we see that preconditioner GS scheme proceeds at an accelerated rate of convergence,
even after the fifth iteration. So for this toy example, the NN preconditioner performs
quiet well.

Finally, a word on what we learn. Since we only take small, non-zero vectors, we update mostly
the bias term and not the weight matrix. We can verify this in the code:

julia> (W_arr[:, :, end] - W_arr[:, :, 1])
16×4 Array{Float64,2}:
  9.2298e-7    3.56583e-6    9.63813e-6   -3.62526e-6
  7.02612e-6   1.83826e-6    4.87793e-6    1.06628e-5
 -1.18359e-5  -5.52383e-6   -1.00905e-5   -2.49149e-5
  4.2548e-6   -4.93141e-7   -5.37068e-6    1.46914e-5
 -4.80919e-6  -1.97348e-5   -5.14656e-5    1.82806e-5
 -2.99721e-5  -1.85244e-5   -4.37366e-5   -3.33036e-5
  5.28254e-5   5.17799e-5    0.000115429   8.27112e-5
 -1.70783e-5  -1.11531e-5   -8.03195e-6   -5.73307e-5
  6.09361e-6   3.89891e-5    8.93109e-5   -3.54379e-5
  5.58404e-5   5.60942e-5    0.000112264   6.67546e-5
 -8.3024e-5   -0.000129324  -0.00020506   -0.000139751
  2.75378e-5   3.47865e-5    1.79561e-5    0.000101882
 -2.89687e-6  -2.14357e-5   -4.965e-5      2.24125e-5
 -3.26029e-5  -3.11671e-5   -5.59135e-5   -4.08413e-5
  5.10635e-5   7.48976e-5    0.000101988   9.20182e-5
 -1.64133e-5  -1.88765e-5   -3.41413e-6   -6.0346e-5

The average matrix element of W has changed only little during learning, on average by
about 0.00001. Looking at how much the elements of b have changed during training

julia> (bmat_arr[:, :, end]*bmat_arr[:, :, 1]') - (bmat_arr[:, :, 1]*bmat_arr[:, :, 1]')
4×4 Array{Float64,2}:
 -0.000674214   0.00163171   0.0012703    0.000217951
  0.00631572    0.00148005  -0.0105803    0.00372173
 -0.050246     -0.0453774   -0.0181303   -0.0511274
  0.011839      0.0128648   -0.00398132   0.00899787

we find that they have changed more, on average by a factor of 100 more than entries of W.
But as discussed earlier, including x0 in the preconditioner is a modeling choice which
one does not have to make.

Conclusions

To summarize, we proposed to use a Neural Network to accelerate an iterative solver
for linear systems by acting as a preconditioner matrix. We propose that the weights
of the neural network can be optimized by automatic differentiation in reverse mode.
By putting the solver in the loop, the training and inference steps couple to the
simulation in a very simple way.

One drawback of the method as written here is that we are limiting ourselfes to initial
guesses \(x \sim 0\). This is due to the requirement of the Gauss-Seidel scheme that the
linear system is positive definite. In more practical settings this can be circumvented
by either using different parameters to be passed to \(P\) than \(x\). Alternatively
one can use an iterative solver that doesn’t pose such restrictions, such as Jacobian-Free
Newton Krylov or Conjugate Gradient etc.