Take a look at these Stack Overflow discussions about this topic, https://stackoverflow.com/questions/65998695/how-to-add-a-l1-or-l2-regularization-to-weights-in-pytorch. https://stackoverflow.com/questions/42704283/l1-l2-regularization-in-pytorch

The key takeaways I have from this discussion is L2 regularization on an individual layer isn’t necessary to train an accurate model. The optimizer can handle the regularization. That is really the optimizer’s job to keep the model training under control. I would suggest consider training a model with different optimizers and hyperparameter choices versus trying to add a regularizer on an individual layer. There is probably a reason the concept isn’t in PyTorch.

I removed the regularization API for exactly the reasons @meanderingstream mentioned:

1. It’s not in PyTorch, and it didn’t seem very commonly used in TensorFlow
2. Regularization is a concern of training/optimization and not the model, so I wanted to keep those decoupled

If you really want to replicate this, you might be able to do something like this, first create a custom layer with a stateful output. In training mode, calls to Axon’s predict method will return a map %{prediction: pred, state: state}, so this will forward the output activations so you can use them in your objective function:

def activity_regularizer(input) do
  Axon.layer(&activity_regularizer_impl/2, [input])
end

deftransformp activity_regularizer_impl(input, opts \\ []) do
  opts = keyword!(opts, mode: :train)
  case opts[:mode] do
    :train ->
      %Axon.StatefulOutput{output: input, state: %{"activations" => input}}

    :inference ->
      # activity regularizer not present during training
      input
  end
end

And then wrap any layer you want to regularize in your custom layer like this:

model =
  Axon.input("data")
  |> Axon.dense(100)
  |> Axon.sigmoid()
  |> activity_regularizer()
  |> Axon.dense(30)
  |> Axon.sigmoid()
  |> activity_regularizer()
  |> Axon.dense(2, activation: :softmax)

Now you can create a custom objective function:

  defn objective(predict_fn, model_state, inputs, y_true) do
    %{prediction: y_pred, state: state} = predict_fn.(model_state, inputs)
    # traverse state for maps with keys "activation"
    penalty = get_activation_penalty(state)
    # get loss and apply penalty to loss
    loss(y_true, y_pred) + penalty
  end

  deftransformp get_activation_penalty(state) do
    # state is a nested map, we can reduce over it
    Enum.reduce(state, Nx.tensor(0.0), fn
      # apply l1 penalty to found activations in state map, you'll need to match on specific
      # layers if you want to apply this penalty differently per-layer
      {_key, %{"activation" => activation}}, acc -> l1_penalty(activation) |> Nx.add(acc)
      _other, acc -> acc
    end)
  end

Then you just need to write the optimization step and integrate that with a custom training loop. I know this probably isn’t very straightforward, given you can just call activity_regularizer in Keras and it just works; however, I think adding a similar API just brings too much coupling between model creation/execution and model training. You are better off using weight decay that’s available in the optimization API to achieve the same regularization affect.