# On the future of Flux.destructure and SciML integration

**URL:** https://discourse.julialang.org/t/on-the-future-of-flux-destructure-and-sciml-integration/104760
**Category:** Machine Learning
**Tags:** flux, sciml
**Created:** [October 9, 2023, 5:42pm UTC](https://discourse.julialang.org/t/on-the-future-of-flux-destructure-and-sciml-integration/104760 "2023-10-09T17:42:55Z")
**Posts on this page:** 5
**Page:** 1

<div class="post-metadata">

### Author: ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)
#### Post date: [October 9, 2023, 5:42pm UTC](https://discourse.julialang.org/t/on-the-future-of-flux-destructure-and-sciml-integration/104760/1 "2023-10-09T17:42:55Z")

</div>

Due to some technical limitations regarding Lux and coupling layers ([1](https://github.com/LuxDL/Lux.jl/issues/416), [2](https://github.com/LuxDL/Lux.jl/issues/417), [3](https://github.com/LuxDL/Lux.jl/issues/420)) and the existence of [InvertibleNetworks.jl](https://github.com/slimgroup/InvertibleNetworks.jl/tree/master), I have chosen to use Flux for my designs. Still, my long-term objective would be to use Optimization.jl and other SciML-related modules, which generally expect flat parameter vectors. Thus, I’ll have to rely on `destructure` and `Restructure`.

The SciML documentation recommends not using Flux, due to some issues related to performance and type promotion ([Flux.jl-vs-Lux.jl](https://docs.sciml.ai/NeuralPDE/stable/#Flux.jl-vs-Lux.jl)). My question: Are these issues the result of some deliberate design choice, arising perhaps from some compromise, or can one hope for a “fix” sometime in the future?

In other words, are these difficulties _intrinsic_ to Flux’s design? If so, is there something one could do to partially bypass them?

Thanks in advance,

---

<div class="post-metadata">

### Author: ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)
#### Post date: [October 9, 2023, 6:42pm UTC](https://discourse.julialang.org/t/on-the-future-of-flux-destructure-and-sciml-integration/104760/2 "2023-10-09T18:42:04Z")

</div>

It’s important to know that these functions had an extremely bad reputation because the original implementation was extremely fragile & had zero tests.

Unfortunately, replacing such a thing means the ire is mostly redirected at the new version (which is much stricter, and much better tested). Perhaps removing it entirely would have been better.

> [@Bizzi](#):
>
> well known that these functions have some issues related to performance and type promotion

This is not really true.

The complaint about type promotion was this. Flux regards a model being Float32 (or Float64 or Float16) as part of the model, and will re-construct it as such. It has explicit functions like `f16` for changing between these. Supplying Float16 parameters to Restructure will not change the model’s precision. Perhaps that’s surprising compared to Base Julia, but GPUs are weird, and accidental Float64 on a GPU was the number one cause of “Flux is much much slower than pytorch” threads. What “well known” means is really that one guy (with a megaphone) was once surprised that the new less-sloppy implementation fixes this. It remains very easy to run models in Float64 if this is actually desired.

Performance should be OK? When last I tried it was fairly easy to cook up benchmarks in which either `destructure` or ComponentArrays.jl is a bit faster – they do similar things but have slight differences in when they allocate copies. Usually neither is a large slice of the model’s run time. If you have particular performance concerns they can probably be solved – there’s no intrinsic reason for large differences.

---

<div class="post-metadata">

### Author: ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)
#### Post date: [October 9, 2023, 7:40pm UTC](https://discourse.julialang.org/t/on-the-future-of-flux-destructure-and-sciml-integration/104760/3 "2023-10-09T19:40:21Z")

</div>

While I can’t really comment on what would be the most natural way to deal with types, I absolutely agree that `destructure` is a great tool, and gets the job done. It also seems to interact much more nicely with tied parameters, which is a major reason I’ll use Flux this time.

However, in regimes where de/restructuring is very frequent (mostly as a hack to use Flux with explicit parameters) the performance cost does seem quite severe, unless I’m doing something wrong. On my machine, small PINNs seem to be 3-4x slower when destructured:

```julia
using Flux, Zygote, BenchmarkTools

C = Chain(Dense(1=>10,tanh),Dense(10=>10,tanh),Dense(10=>10,tanh),Dense(10=>1))
ps, re = Flux.destructure(C)
v = [1f0]

@btime C(v)
@btime re(ps)(v)
@btime Zygote.gradient(c->c(v)[1],C)
@btime Zygote.gradient(p->re(p)(v)[1],ps)

```

This is, again, probably stretching the tool beyond what it was made for; I’m glad it exists in any case. My only question is if there’s anything I can do to reduce the slowdown - like designing custom layers or avoiding deep nesting of parameters.

(In any case, “well known” was a poor choice of words, and I have edited the OP to reflect it)

---

<div class="post-metadata">

### Author: ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)
#### Post date: [October 9, 2023, 7:55pm UTC](https://discourse.julialang.org/t/on-the-future-of-flux-destructure-and-sciml-integration/104760/4 "2023-10-09T19:55:33Z")

</div>

> [@Bizzi](#):
>
> where de/restructuring is very frequent (mostly as a hack to use Flux with explicit parameters) the performance cost does seems quite severe, unless I’m doing something wrong. On my machine, small PINNs seem to be 3-4x slower when destructured

In this regime it is expensive. While I haven’t checked this example today, my claim above is that using ComponentArrays.jl will also be expensive. I’d say that explicit/implicit is not the right axis here, it’s more like solid-vector vs. nested structure. For sufficiently small networks, converting back & forth may take longer than the matrix multiplications etc.

There is more that could be done here, e.g. Optimisers.jl could add a `destructure!` which re-uses the same model quite easily. It might save quite a bit here. See e.g. [issue 146](https://github.com/FluxML/Optimisers.jl/issues/146).

This is the regime where [SimpleChains.jl](https://github.com/PumasAI/SimpleChains.jl) is likely to be much faster. It’s a completely different design, which never uses a nested structure at all. But has various other limitations.

> [@Bizzi](#):
>
> more nicely with tied parameters

This is a key difference from ComponentArrays.jl, which perhaps we should highlight. (As well as being the main source of coding headaches!) I’m glad if it’s useful to someone.

```julia
julia> using ComponentArrays, Optimisers

julia> let twice = [1.0, 2.0]
        cv = ComponentArray(x=twice, y=twice, z=[1.0, 2.0])
        cv.x[1] += 999
        cv # this has 6 independent scalar parameters
       end
ComponentVector{Float64}(x = [1000.0, 2.0], y = [1.0, 2.0], z = [1.0, 2.0])

julia> let twice = [1.0, 2.0]
        v, re = destructure((x=twice, y=twice, z=[1.0, 2.0]))
        @show v # only 4 indep parameters
        v[1] += 999
        re(v)
       end
v = [1.0, 2.0, 1.0, 2.0]
(x = [1000.0, 2.0], y = [1000.0, 2.0], z = [1.0, 2.0])

```

---

<div class="post-metadata">

### Author: ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)
#### Post date: [October 9, 2023, 8:24pm UTC](https://discourse.julialang.org/t/on-the-future-of-flux-destructure-and-sciml-integration/104760/5 "2023-10-09T20:24:50Z")

</div>

Fantastic, thanks. `re` making a view instead of making a copy seems like a relatively straightforward fix; I’ll try and get it running on my end.
