diff --git a/lib/axon.ex b/lib/axon.ex index 764dd56b..e2e5eda3 100644 --- a/lib/axon.ex +++ b/lib/axon.ex @@ -3425,6 +3425,47 @@ defmodule Axon do layer(:bias, [x, bias], name: opts[:name], meta: opts[:meta], op_name: :bias) end + @doc """ + Adds a learnable per-channel scale layer to the network. + + A scale layer multiplies the input elementwise by a trainable + parameter aligned to the input's `:channel_index` axis. + + This is commonly used as the `gamma` in residual blocks of modern + Transformer architectures (CaiT, ConvNeXt, BEiT, EVA, etc.). + + ## Options + + * `:name` - layer name. + + * `:scale_initializer` - initializer for the scale weights. + Defaults to `Axon.Initializers.full(1.0e-6)`. + + * `:channel_index` - input feature axis along which the scale is + broadcast. Defaults to `-1`. + """ + @doc type: :linear + def scale(%Axon{} = x, opts \\ []) do + opts = + Keyword.validate!(opts, [ + :name, + :meta, + scale_initializer: Axon.Initializers.full(1.0e-6), + channel_index: -1 + ]) + + channel_index = opts[:channel_index] + + scale = param("scale", [{:axis, channel_index}], initializer: opts[:scale_initializer]) + + layer(:scale, [x, scale], + name: opts[:name], + meta: opts[:meta], + channel_index: channel_index, + op_name: :scale + ) + end + @doc """ Adds a stack columns layer to the network. diff --git a/lib/axon/layers.ex b/lib/axon/layers.ex index bc4b0582..bac0aab3 100644 --- a/lib/axon/layers.ex +++ b/lib/axon/layers.ex @@ -2001,6 +2001,32 @@ defmodule Axon.Layers do input + bias end + @doc ~S""" + Functional implementation of a learnable scale layer. + + Multiplies the input elementwise by `scale` broadcast along the + `:channel_index` axis: + + $$y = x * \gamma$$ + + where $\gamma$ is the scale parameter with one element per + position along `:channel_index`. + + ## Options + + * `:channel_index` - input axis along which the scale is + applied. Defaults to `-1`. + """ + defn scale(input, scale, opts \\ []) do + opts = keyword!(opts, channel_index: -1, mode: :train) + + num_channels = Nx.axis_size(input, opts[:channel_index]) + parameter_shape = norm_parameter_reshape(input, num_channels, opts[:channel_index]) + + scale = Nx.reshape(scale, parameter_shape) + Nx.multiply(input, scale) + end + @doc """ Resizes a batch of tensors to the given shape using one of a number of sampling methods. diff --git a/test/axon/compiler_test.exs b/test/axon/compiler_test.exs index c152b114..d946d680 100644 --- a/test/axon/compiler_test.exs +++ b/test/axon/compiler_test.exs @@ -347,6 +347,70 @@ defmodule CompilerTest do end end + describe "scale" do + test "initializes in default case" do + model = Axon.input("input_0", shape: {nil, 3}) |> Axon.scale(name: "scale") + + input = random({1, 3}) + + assert {init_fn, _predict_fn} = Axon.build(model) + + assert %ModelState{ + data: %{"scale" => %{"scale" => scale}}, + parameters: %{"scale" => ["scale"]} + } = init_fn.(input, ModelState.empty()) + + assert Nx.shape(scale) == {3} + assert Nx.type(scale) == {:f, 32} + end + + test "applies small init multiplicatively" do + model = + Axon.input("input_0", shape: {nil, 3}) + |> Axon.scale(name: "scale", scale_initializer: Axon.Initializers.full(1.0e-6)) + + input = Nx.tensor([[1.0, 2.0, 3.0]]) + + assert {init_fn, predict_fn} = Axon.build(model) + params = init_fn.(input, ModelState.empty()) + + out = predict_fn.(params, input) + assert_all_close(out, Nx.tensor([[1.0e-6, 2.0e-6, 3.0e-6]]), atol: 1.0e-12) + end + + test "broadcasts correctly along a non-trailing channel_index" do + model = + Axon.input("input_0", shape: {nil, 2, 4}) + |> Axon.scale(name: "scale", channel_index: 1) + + input = + Nx.tensor([ + [ + [1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0] + ] + ]) + + assert {init_fn, predict_fn} = Axon.build(model) + params = init_fn.(input, ModelState.empty()) + + params = + Axon.ModelState.update(params, %{"scale" => %{"scale" => Nx.tensor([2.0, 3.0])}}) + + out = predict_fn.(params, input) + + expected = + Nx.tensor([ + [ + [2.0, 2.0, 2.0, 2.0], + [3.0, 3.0, 3.0, 3.0] + ] + ]) + + assert_equal(out, expected) + end + end + describe "dense" do test "initializes in default case" do model = Axon.input("input_0", shape: {nil, 1}) |> Axon.dense(1, name: "dense")