Moving @compact to Flux?

[MilesCranmer]

Hi all,

I was wondering when @compact could be moved to Flux.jl?

My experience with it has been great thus far. Every single time I make a neural net in Julia I use a mix of Chain for simple linear networks and @compact for sticking operations and parameters together.

I’ve shown a few people it when introducing them them to Flux.jl and they also find it very intuitive. Would be great to include in Flux.jl to help accelerate user growth!

Cheers,
Miles

[ToucheSir]

Hey @MilesCranmer, please bear with us since this is literally the first time we need to consider how to do such a move. I'm not sure how everyone feels about it, so ping @FluxML/committers to get some input. If you want to discuss this in a sync fashion, we also have a ML call this Friday (see the Julia events calendar)—let me know if you have time and are interested, and I'll add it to the agenda.

[ToucheSir]

https://github.com/FluxML/FluxML-Community-Call-Minutes/blob/main/meeting-minutes/2023-07-14.md

[MilesCranmer]

Hi @ToucheSir,

So sorry for the late notice but I'm unable to attend today, an important work meeting came up and is overlapping with the time. Hopefully I can join on a future call. But please go ahead and feel free to discuss the @compact integration in my absence, I think you and @mcabbott are knowledgeable about how it works and the motivations for it.

Best wishes,
Miles

[darsnack]

We spoke about this on the call. There were no major objections to including it in Flux.jl. Would you have time to make a PR? One note made during the call is that we want the docstring to mention this should be used for simple layers, and that more complex use cases might be better served by a struct.

[MilesCranmer]

Would you have time to make a PR?

Unfortunately time is quite scarce for me these days as I'm helping create a new masters program for this fall. But I will try...

One note made during the call is that we want the docstring to mention this should be used for simple layers, and that more complex use cases might be better served by a struct.

Fine with me!

[darsnack]

Unfortunately time is quite scarce for me these days as I'm helping create a new masters program for this fall. But I will try...

No worries if you are bandwidth limited right now. One of us can handle it too.

[mcabbott]

I have strong reservations about copying this across. It seems like a first draft, and maybe in its present state very useful for some task / some ways of thinking. This is exactly what Fluxperimental is for. It could be merged here, and available as a registered package, without waiting to figure out the ideal thing. (Nor worrying about growing Flux's docs to have multiple ways of doing the same thing. Let alone worrying about testing & documenting as many edge cases as possible.)

But the possible space of slightly similar designs is large. Are we sure there isn't some way to have 90% the benefit (or 200%) for 5% the code / complexity? E.g. if we kill @functor maybe just writing a let block becomes a great way to work, with no clever macros at all?

What I think would be very helpful is examples of actual use. The documentation just re-makes things we can already do; we don't want to replace struct Dense with this. Are there nice models which work well this way, and don't in other ways? Can e.g. something in the model zoo be made much cleaner like this?

[MilesCranmer]

For example, in the MHA snippet, each head is its own type.

This is not true, Julia is able to figure out it is the same function definition, and inline it:

julia> typeof(multihead_self_attention(10, 5, 32, 32, 3))
Fluxperimental.CompactLayer{var"#13#17", NamedTuple{(:out, :heads), Tuple{String, String}}, NamedTuple{(:out, :heads), Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Vector{Fluxperimental.CompactLayer{var"#12#16"{Int64}, NamedTuple{(:K, :V, :Q), Tuple{String, String, String}}, NamedTuple{(:K, :V, :Q), Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}}}}}

You can see each head has the function var"#12#16". So it's a single type.

Another example just out of personal curiosity, now with different # of parameters in each head. You can get the layer type from `typeof(head(...))`. Not sure when you would use this, but it seems to work.

julia> function head(num_features; d_attn=512, d_value=512)
           @compact(
               K = Dense(num_features => d_attn),
               V = Dense(num_features => d_value),
               Q = Dense(num_features => d_attn)
           ) do x
               k, v, q = K(x), V(x), Q(x)
               x = sum(k .* q; dims=1) ./ sqrt(d_attn)
               softmax(x; dims=2) .* v
           end
       end

julia> typeof(head(32))
Fluxperimental.CompactLayer{var"#20#21"{Int64}, NamedTuple{(:K, :V, :Q), Tuple{String, String, String}}, NamedTuple{(:K, :V, :Q), Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}

julia> typeof(head(50))
Fluxperimental.CompactLayer{var"#20#21"{Int64}, NamedTuple{(:K, :V, :Q), Tuple{String, String, String}}, NamedTuple{(:K, :V, :Q), Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}

[darsnack]

This is not true, Julia is able to figure out it is the same function definition

You're right, the type is only defined once when we define multihead_self_attention's body. Though if you modify multihead_self_attention, the type changes, so it still isn't something stable that you can use for dispatch.

The modification time is what I thought was the point of comparison.

Fair, so basically we need ProtoStruct.jl working. RedefStructs.jl does not work in Julia 1.8+.

An alternative solution would be to have @compact be included in Flux, but we don't introduce it in the normal places in the docs for defining custom layers. Instead, it can be pitched as a prototyping utility with the recommendation to move to an explicit struct once you are ready.

[MilesCranmer]

You're right, the type is only defined once when we define multihead_self_attention's body. Though if you modify multihead_self_attention, the type changes, so it still isn't something stable that you can use for dispatch.

Agree!

An alternative solution would be to have @compact be included in Flux, but we don't introduce it in the normal places in the docs for defining custom layers. Instead, it can be pitched as a prototyping utility with the recommendation to move to an explicit struct once you are ready.

Sounds like a good compromise to me.

[mcabbott]

Wait how's this a compromise, on the question of whether or not to copy this code across?

but we don't introduce it in the normal places in the docs for defining custom layers. Instead, it can be pitched as a prototyping utility with the recommendation to move to an explicit struct once you are ready.

This is a documentation project which could be tacked right now. Maybe it sits alongside http://fluxml.ai/Flux.jl/stable/models/advanced/ or maybe better, it replaces that with something not-confusing, which ends up discussing this macro. There is no rule that Flux's "tutorials" section not document things from this package.

[MilesCranmer]

I guess the compromise was:

+ @compact be included in Flux

but

- emphasize it only as a prototyping tool

[MilesCranmer]

I don't want to derail things with new feature ideas as I think current @compact is good to go and we are starting to find a compromise, but I just want to leave a future idea here, regarding this point from @darsnack:

You're right, the type is only defined once when we define multihead_self_attention's body. Though if you modify multihead_self_attention, the type changes, so it still isn't something stable that you can use for dispatch.

Here's an option for getting around it:

Rewrite CompactLayer to be parametrized by a symbol:

struct CompactLayer{S,F,NT1<:NamedTuple,NT2<:NamedTuple}
  symbol::Val{S}
  fun::F
  name::Union{String,Nothing}
  strings::NTuple{3,String}
  setup_strings::NT1
  variables::NT2
end
CompactLayer(f::Function, name::Union{String,Nothing}, str::Tuple, setup_str::NamedTuple, symb=:Default; kw...) = CompactLayer(Val(symb), f, name, str, setup_str, NamedTuple(kw))

which is by default, just Val{:Default}.

This would be used like:

@compact(Linear, W=randn(5, 5), b=randn(5)) do x
    W * x .+ b
end

so that you can dispatch on specific layers using their symbol:

function something(m::CompactLayer{:Linear})
    ...
end

which would be stable, as it is insensitive to the specific function used in m, so long as it has the correct symbol. So this would be completely compatible with prototyping new layers.

---

Edit: made a PR #14

[mcabbott]

Call it trying to home in on a minimal solution, rather than walking in circles!

[MilesCranmer]

These @proto methods are indeed nicer than the current strategy. But it takes three separate blocks to modify the method: one for the struct, one for the method, and one for the initialization. If we emphasize that @proto is just for prototyping things, one single block (and no need to specify types) for all three would be the nicest, right?

So the next logical conclusion is to have something like:

@proto struct Linear
    W=randn(32, 32)
    b=randn(32)

    (x) -> W * x + b
end
model = Linear()

but this isn't really a struct, it's more like a let. So we end up at

model = @proto let W=randn(32, 32), b=randn(32)
    function (x)
        W * x + b
    end
end

but we don't need that inner block, so we simplify to

@proto(W=randn(32, 32), b=randn(32)) do x
    W * x + b
end

In other words I feel like we might already be at a local minima in design space?

[mcabbott]

isn't really a struct

Well, it has a name. Constructors use this name. It also has to be defined at top level, probably.

@compact has some name field too, but no examples use it so far. The only example above defines a constructor with a different name, multihead_self_attention.

Maybe there's other complexity lying around. For example in @compact(n = 10, W=randn(n,n), b=randn(n)) do x what exactly is n, what scope rules does it obey, how do readers know this? This would not be a legal set of keyword arguments, they are more like statements, but not quite. I think I am to blame for this one... late one night for fun, trying to be maximal, really.

[mcabbott]

it takes three separate blocks to modify the method: one for the struct, one for the method, and one for the initialization

Why is that bad?

Understanding that a Flux model has these three things is important. Storage, forward pass, initialisation. Knowing which bits run once, and don't need to be Zygote-friendly, is important.

Maybe this is part of my complaint about @compact, that it seems to obscure this.

[MilesCranmer]

I think at this point (because of an upcoming period of travel/other craziness) I need to step away and let you three sort it out. I'm not sure what else I can contribute mostly because I lack an understanding of the motivations for the existing API. It's becoming a bit of a time sink to keep discussing it, though I agree it's important. Happy to give an opinion here and there though.

[mcabbott]

Here's a made-up example of @compact which makes a layer containing parameters directly. It does some preliminary work to initialise. And this initialisation is a bit weird in that:

• the syntax is neither that of code nor that of keyword arguments
• intermediate steps are all stored, and
• it all gets printed every time, even after training has altered the parameters.

julia> using Fluxperimental: @compact; using Flux

julia> mylayer(n) = @compact(  # this layer constructor is called mylayer
           m = 2n,
           tmp = randn32(m, n),  # initialisation... these are not legal keyword arguments
           mat = tmp ./= sum(tmp; dims=1)  # ... and need not be Zygote-friendly
       ) do x
           y = mat * x
           y[1:n, :] .+ relu(y[n+1:end, :])   # forward pass 
       end
mylayer (generic function with 1 method)

julia> m = Chain(mylayer(3), Dense(3 => 1))  # notice that this stores tmp, and prints the init code
Chain(
  @compact(
    m = 2n,
    tmp = randn32(m, n),                # 18 parameters
    mat = tmp ./= sum(tmp; dims = 1),   # 18 parameters
  ) do x 
      y = mat * x
      y[1:n, :] .+ relu(y[n + 1:end, :])
  end,
  Dense(3 => 1),                        # 4 parameters
)         # Total: 3 trainable arrays, 22 parameters,
          # plus 1 non-trainable, 18 parameters, summarysize 754 bytes.

julia> m[1].variables
(m = 6, tmp = Float32[0.65733784 0.332404 0.06522033; 0.40533376 -0.18564595 0.38826385; … ; 0.17651373 0.089108355 0.22100908; -0.08086931 1.1042991 -0.32503128], mat = Float32[0.65733784 0.332404 0.06522033; 0.40533376 -0.18564595 0.38826385; … ; 0.17651373 0.089108355 0.22100908; -0.08086931 1.1042991 -0.32503128])

(Is there really a non-trainable parameter array?)

We got to this design starting from one which stored a NamedTuple. Besides being one line, this has the advantage that init obvioulsy runs just once, and is obviously normal Julia code.

The downsides are that

• you have to qualify fields like store.mat
• printing doesn't remind you about where this came from.

julia> layerfactory(f; layers...) = Base.Fix1(f, NamedTuple(layers));

julia> function mylayer1(n)
           (_, m, _) = n, 2n, 3n  # useless destructuring, would be illegal in macro
           tmp = randn32(m, n)    # tmp will be discarded after this runs
           mat = tmp ./ sum(tmp; dims=1)

           return layerfactory(; mat) do store, x
               y = store.mat * x   # forward pass needs to qualify fields
               y[1:n, :] .+ relu(y[n+1:end, :])
           end
       end;

julia> m1 = Chain(mylayer1(3), Dense(3 => 1))  # clearly the printing of Fix1 isn't ideal
Chain(
  Fix1(
    var"#37#38"{Int64, Matrix{Float32}}(3, Float32[0.16561371 0.6836739 -19.871439; 0.30575347 0.19557928 -14.132152; … ; 0.2924494 -0.20139198 36.675316; 2.0378335 -0.033085346 -5.4660473]),
    (mat = Float32[0.16561371 0.6836739 -19.871439; 0.30575347 0.19557928 -14.132152; … ; 0.2924494 -0.20139198 36.675316; 2.0378335 -0.033085346 -5.4660473],),  # 18 parameters
  ),
  Dense(3 => 1),                        # 4 parameters
)                   # Total: 3 arrays, 22 parameters, 352 bytes.


julia> layerfactory(a=Dense(1=>2), b=Dense(1=>2)) do nt, x  # printing is better here, not ideal
        nt.a(x) + nt.b(x)
       end |> Chain
Chain(
  Fix1(
    var"#43#44"(),
    NamedTuple(
      Dense(1 => 2),                    # 4 parameters
      Dense(1 => 2),                    # 4 parameters
    ),
  ),
)                   # Total: 4 arrays, 8 parameters, 352 bytes.

Can the problems with the 1-line way be addressed by something simpler?

I keep asking for examples because it's not clear to me whether this kind of layer is intended to be in-scope, or not. If not, then maybe init steps like this ought to be forbdden? Maybe the contents must always be other layers, never arrays?

[mcabbott]

flux printing is supposed to print a code block you can copy-paste and execute to re-init

Maybe. It is strongly broken by the above examples, as variables go out of scope -- Guarav's point 2. It's more weakly broken by e.g. Dense(3 => 4; init=Flux.identity_init).

custom macro to work with $ interpolation

Sure, it's just e.g. Expr(:$, 1). But I'm not sure I see what the proposed behaviour is.

making issue(s) on this repo for the strange behavior

For some reason we decided not to have issues, just discussion.

[gaurav-arya]

I'd advocate for having issues, and making concrete issues for each class of problems noted in this thread. At least, I feel like disallowing bug reports here is at odds with the philosophy that the features here should be battle-tested prior to the PR to Flux.

[MilesCranmer]

Current behavior:

julia> function model(n, m)
           @compact(f=Dense(n=>m), α=rand(m)) do x
               f(x) .+ cos(α)
           end
       end
model (generic function with 1 method)

julia> model(3, 4)
@compact(
  f = Dense(3 => 4),                    # 16 parameters
  α = rand(m),                          # 4 parameters
) do x
    f(x) .+ cos(α)
end                  # Total: 3 arrays, 20 parameters, 466 bytes.

proposed behavior would be to allow $ to interpolate when printing:

julia> function model(n, m)
           @compact(f=Dense(n=>m), α=rand($m)) do x
               f(x) .+ cos(α)
           end
       end
model (generic function with 1 method)

julia> model(3, 4)
@compact(
  f = Dense(3 => 4),                    # 16 parameters
  α = rand(4),                          # 4 parameters
) do x
    f(x) .+ cos(α)
end                  # Total: 3 arrays, 20 parameters, 466 bytes

or even $(rand(m)) to get the values themselves.

This way, users have some control over displaying models.

[mcabbott]

But should that be f=Dense($n=>$m), α=rand($m) for consistency? When would the no-$ case (preserving the name of undefined variables) be desirable anyway?

The present printing is a strange hybrid really... I'm not even sure how it's decided that sometimes the expression is printed, and sometimes the result:

julia> let m=2, b=false
         @compact(f=Dense(diagm(ones(m^2))), α=rand(m^2)./3, g=Flux.trainmode!(Dropout(1/m), b)) do x
           f(x) .+ cos.(α) |> g
         end
       end
@compact(
  f = Dense(4 => 4),                    # 20 parameters
  α = rand(m ^ 2) ./ 3,                 # 4 parameters
  g = Dropout(0.5, active=false),
) do x 
    f(x) .+ cos.(α) |> g
end                  # Total: 3 arrays, 24 parameters, 680 bytes.

[MilesCranmer]

Yeah it calls the regular _big_show method if it's a layer. But we could change that for sure:

Fluxperimental.jl/src/compact.jl

Lines 172 to 178 in 9dcae27

	if Flux._show_leaflike(v)
	# If the value is a leaf, just print verbatim what the user wrote:
	str = String(k) * " = " * setup_strings[k]
	_just_show_params(io, str, v, indent+2)
	else
	Flux._big_show(io, v, indent+2, String(k))
	end

[MilesCranmer]

Ping on this. Would be really nice to have!

[gaurav-arya]

Hey, I'd be interested in helping push this through for fun :) Regarding @mcabbott's point above -- do you think it would make sense for the initialization steps to be independent rather than sequential, so that they closely resemble kwarg initialization? If so, I'd be interested in making a PR to implement that and document it as such.

[gaurav-arya]

Airing it out for any comments before I do it: I hope to replace the recursive addition of the self. prefix, which I worry might be a bit brittle (e.g. what if w is a param and user declares an anonymous function w->w^2 inside the do block?) with a let block added by the macro the wraps the user-provided function like so(in phone written pseudocode...):

(self, x) -> let a = self.a, W = self.W, ... begin
     do x -> ... end # this is the user written function 
end

I think this should get all the same functionality. And means that minus the printing logic, @compact should become an elegant combination of two ideas that have been floated around in this discussion: NamedTuples (basically what CompactLayer is) and let blocks. I think that's a point worth reiterating: the @compact macro would not be its own magic solution orthogonal to alternatives, but rather some pixie dust lightly sprinkled over the alternative pure Julia solutions that have been discussed. Indeed, the @macroexpand, if we squint hard enough so that the printing logic fades away, should basically be a let-block-wrapped user-provided function provided to a named-tuple-backed prototype-friendly callable struct.

Let me know if this makes sense!

[gaurav-arya]

Oh yup, you're right about it just being the function body.

Edit: actually, in the metaprogramming I should probably drop the x outside so I can just paste in the user function without separating args. Then it's easy enough to uncurry it, (self, args...) -> f(self)(args...)

[MilesCranmer]

@gaurav-arya see FluxML/Flux.jl#2107 (comment) for an earlier approach to @compact

struct LayerFactory{F<:Function,NT<:NamedTuple}
    forward::F
    layers::NT
end

LayerFactory(f; layers...) = LayerFactory(f, NamedTuple(layers))

function (f::LayerFactory)(args...)
    return f.forward(f.layers, args...)
end

@functor LayerFactory

used like

my_layer = LayerFactory(; w1=Dense(5, 128), act=relu) do self, x
    self.act(self.w1(x))
end

The reason for the macro was to print things like randn(3) exactly as the user wrote, and also avoid the need for self.. But in all honesty I’d be happy with either one.

Maybe to make the existing macro more robust we could just have a less common variable name than self...?

[gaurav-arya]

Hey, regarding the automatic self., I just put in #17 which should hopefully be a pretty a robust solution:) [And also an idea that has precedent: it's just a let block].

Subject to review on #16 and #17, the last thing to sort out seems to be the printing. For convenience printing, my feeling is that we should concretize by default rather than do pure-source-code by default. E.g. if the user writes

function f(n)
    my_special_func(x) = x^2
    @compact(w = my_special_func(10)) do x
    end
end

we should probably print 100 for w rather than the out-of-scope my_special_func. We can catch special cases though, where it is more enlightening to concretize only the arguments and not the function, the prominent examples being Flux layers, randn and friends, etc. The nice thing about printing is that we don't have to get it perfect, and can fall back on concretizing arguments whenever we want, improving over time in a non-breaking way.

You're exactly write that automatic self. and the ability to do convenience printing are the two things the macro provides. If we think these two things are worth it over a LayerFactory, I'm happy to try and work on tweaking the printing to follow the above philosophy.. (Maybe they aren't worth it? It deserves some thought, opinions welcome from you, @mcabbott, and others. But with #16, #17, and a printing solution like above, it does seem like the technical challenges to @compact are solvable/solved, so the decision can at least be made with the knowledge that @compact can be implemented in a reasonable way).

[mcabbott]

One thought is that with #17 we could store the lambda you provide (instead of storing a string, but re-writing the code) [edit] alongside a function to un-wrap. That means we could consider outsourcing more to https://github.com/MasonProtter/LegibleLambdas.jl , instead of handling the printing of the forward pass here. Could give something that's easier to inspect?

[gaurav-arya]

LegibleLambdas looks like a good option for handling the function call printing. It matches our heuristic here in the most common cases, but does a bit better in the cases where the function closes over variables, so it's very appealing.

Would LegibleLambdas.jl be feasibly accepted as a downstream dependency of Flux? It's about a 100 lines of very simple code, with no bug reports, and a single dependency on MacroTools. (If not, I'd rather replicate the code here and/or settle for a less perfect heuristic/compromise.)