TF2-style layer losses

[drdozer]

Is it possible to declare layer-associated losses, as per tensorflow? In the TF2 API, layers can be declared with optional weight and activation losses. It also provides some standard loss functions including L1 and L2. During training, as well as collecting the parameters, it collects the losses. Perhaps there would be some way to retro-fit this to flux? Perhaps using the functor API to walk the graph, and collect losses if that layer provides them?

[ToucheSir]

Do you mean "regularizers" as in https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense#args? Currently Flux handles that much like PyTorch, by integrating common regularizers as optimization steps (e.g. WeightDecay). This avoids the overhead of having to AD through said regularization. Of course, you can still just calculate per-parameter values by iterating over them in gradient. Flux.modules was designed to help with this.

Speaking more philosophically, I and likely many other contributors are of the opinion that (tf.)Keras crams too much orthogonal functionality into its layer classes. Thankfully, we are not bound by the same constraints or legacy that TF is, so no retrofitting is required :)

[darsnack]

This may be of relevance as well: #1575

[DhairyaLGandhi]

A simple answer is to have a function that accepts a layer and performs regularisation as the forward pass. If you want to adapt it to specific layers, you'd wrap reg_loss(m) = x -> m(x) .+ sum(sqnorm, Flux.params(m)) and pass it to a chain as Chain(reg_loss(Dense(3,3))).

Easiest is to pass the params you want regularised to an l2 loss function that can be added to the loss. loss + sum(sqnorm, Flux.params(layers))

[ToucheSir]

Since the Slack thread mentioned wanting something more "declarative" (which we can go back and forth on, but I'm reusing the phrasing), you could envision extracting the reg_loss function Dhairya posted into a callable struct that wraps a layer, forwards through arguments and implements the functor interface to extract the sublayer for params. The call method of this struct could then be parameterized in terms of the wrapped layer via multiple dispatch, with a fallback to Functors.children or Flux.trainable to extract relevant params for regularization. A further extension would be to also parameterize the struct in terms of the regularization function itself. Note that at no point did this require any change to Flux itself—no parameter explosion or modification of existing layers required!

[darsnack]

This design does just that (it adds a gradient hook but you could do a loss term instead too).

[drdozer]

Yes, I mean regularizers. Thanks for clarifying that @ToucheSir. If regularization is manually described within the loss function, then the loss function now potentially needs to white-box the network it is calculating the loss for. This places knowledge about the network design in two places that have to be updated in parallel, particularly if different parameters are under different regularization regimes.

I agree that blindly mimicking the T2 API is not a good aspiration to have.

How would we calculate a regularizer loss on the activations of a layer rather than weights?

[darsnack]

How would we calculate a regularizer loss on the activations of a layer rather than weights?

You could use Flux.activations to get the layer-by-layer activations of a Chain then use that output in your loss function.

If your regularizer can be implemented as a gradient hook, then the approach I linked above would also work.

[drdozer]

@darsnack I'm a julia noob - I'm not really clear about what a gradient hook is. But I think perhaps this is illustrating the issue. Should I, as a noob, need to know about gradient hooks to be able to tell flux to apply an L1 regularization on the biases and L2 regularization on the weights of the first dense layer in my multi-layer network? It seems like a big asymmetry of effort, compared to the simplicity of setting up a chain describing a linear pipeline of layers.

[darsnack]

A gradient hook is an addition to a function/layer that does nothing on the forward pass but modifies the weight updates on the backward pass. Many regularizers can be expressed analytically in this way. Without looking at the TF2 source, I bet that this is their approach too.

As for whether you should be aware of it, absolutely not! A user shouldn't have to manually implement these features. I'm just offering possible design options for the implementation. If you absolutely need "declarative" regularizers, you can copy the code. But it's on the Flux team to settle on a design and implement it.

[DhairyaLGandhi]

We usually try to keep the API minimal and encourage the use of regular Julia to express these things esp since they are very minimal in terms of code reuse. I agree it's not incumbent on the user to know how to implement standard regularisation, but we wouldn't want to enforce a blessed pattern to implement new ones.

[ToucheSir]

Should I, as a noob, need to know about gradient hooks to be able to tell flux to apply an L1 regularization on the biases and L2 regularization on the weights of the first dense layer in my multi-layer network?

No, for surgical application of parameter norms (or any function over layer params), you can write them out in the forward pass as with any other framework:

# l1norm(x) = ... , l2norm(x) = ...
loss(m(x), y) + l1norm(m[1].bias) + l2norm(m[1].weight)

This could be easily extracted into a custom function that dispatches on different layer types as applicable. Combo that with modules to work over multiple layers and you're good to go :)