[FEATURE] Integrate gating & blending activations into WaveNet (#193)

* Replaced manual gating/blending with Gating/BlendingActivation classes.

* Deleted debugging stuff and replaced throw with assert.

* Remove backward compatibility constructors from _Layer, LayerArrayParams, and _LayerArray

- Removed boolean gated parameter constructors from _Layer class
- Removed boolean gated parameter constructors from LayerArrayParams struct
- Removed boolean gated parameter constructors from _LayerArray class
- Updated all test files to use GatingMode enum instead of bool gated
- Backward compatibility logic retained in Factory function for JSON parsing
- All tests pass after updating to use GatingMode enum

* Formatting

* Enhance _Layer constructor to validate gating and blending activations

- Added validation checks in the _Layer constructor to ensure that gating and blending activations are not provided when the gating mode is set to NONE.
- Updated the constructor signature to remove default values for gating and blending activations, requiring explicit values during instantiation.
- Adjusted related test cases to reflect the new constructor requirements, ensuring comprehensive coverage of the updated functionality.
- Improved error handling by throwing std::invalid_argument exceptions for invalid activation configurations.

* Refactor _Layer constructor to use secondary_activation parameter

- Updated the _Layer constructor to replace gating and blending activation parameters with a single secondary_activation parameter.
- Added validation to ensure secondary_activation is provided for GATED and BLENDED modes.
- Adjusted related test cases to reflect the new constructor signature and ensure proper functionality.
- Removed default values for gating and blending activations, enforcing explicit activation configuration during instantiation.

* Note on gating/blending activations

---------

Co-authored-by: João Felipe Santos <santosjf@pm.me>

Author: sdatkinson

NAM/gating_activations.h

@@ -37,7 +37,13 @@ class GatingActivation
   , gating_activation(gating_act)
   , num_channels(input_channels)
   {
-    assert(num_channels > 0);
+    if (num_channels <= 0)
+    {
+      throw std::invalid_argument("GatingActivation: number of input channels must be positive");
+    }
+    // Initialize input buffer with correct size
+    // Note: current code copies column-by-column so we only need (num_channels, 1)
+    input_buffer.resize(num_channels, 1);
   }
 
   ~GatingActivation() = default;
@@ -47,7 +53,8 @@ class GatingActivation
    * @param input Input matrix with shape (input_channels + gating_channels) x num_samples
    * @param output Output matrix with shape input_channels x num_samples
    */
-  void apply(Eigen::MatrixXf& input, Eigen::MatrixXf& output)
+  template <typename InputDerived, typename OutputDerived>
+  void apply(const Eigen::MatrixBase<InputDerived>& input, Eigen::MatrixBase<OutputDerived>& output)
   {
     // Validate input dimensions (assert for real-time performance)
     const int total_channels = 2 * num_channels;
@@ -59,6 +66,9 @@ class GatingActivation
     const int num_samples = input.cols();
     for (int i = 0; i < num_samples; i++)
     {
+      // Store pre-activation input values in buffer to avoid overwriting issues
+      input_buffer = input.block(0, i, num_channels, 1);
+
       // Apply activation to input channels
       Eigen::MatrixXf input_block = input.block(0, i, num_channels, 1);
       input_activation->apply(input_block);
@@ -87,6 +97,7 @@ class GatingActivation
   activations::Activation* input_activation;
   activations::Activation* gating_activation;
   int num_channels;
+  Eigen::MatrixXf input_buffer;
 };
 
 class BlendingActivation
@@ -103,10 +114,8 @@ class BlendingActivation
   , blending_activation(blend_act)
   , num_channels(input_channels)
   {
-    if (num_channels <= 0)
-    {
-      throw std::invalid_argument("BlendingActivation: number of input channels must be positive");
-    }
+    assert(num_channels > 0);
+
     // Initialize input buffer with correct size
     // Note: current code copies column-by-column so we only need (num_channels, 1)
     input_buffer.resize(num_channels, 1);
@@ -119,7 +128,8 @@ class BlendingActivation
    * @param input Input matrix with shape (input_channels + blend_channels) x num_samples
    * @param output Output matrix with shape input_channels x num_samples
    */
-  void apply(Eigen::MatrixXf& input, Eigen::MatrixXf& output)
+  template <typename InputDerived, typename OutputDerived>
+  void apply(const Eigen::MatrixBase<InputDerived>& input, Eigen::MatrixBase<OutputDerived>& output)
   {
     // Validate input dimensions (assert for real-time performance)
     const int total_channels = num_channels * 2; // 2*channels in, channels out

NAM/wavenet.cpp

@@ -58,29 +58,38 @@ void nam::wavenet::_Layer::Process(const Eigen::MatrixXf& input, const Eigen::Ma
     this->_conv.GetOutput().leftCols(num_frames) + _input_mixin.GetOutput().leftCols(num_frames);
 
   // Step 2 & 3: activation and 1x1
-  if (!this->_gated)
+  //
+  // A note about the gating/blending activations:
+  // They take 2x dimension as input.
+  // The top channels are for the "primary" activation and will be in-place modified for the final result.
+  // The bottom channels are for the "secondary" activation and should not be used post-activation.
+  if (this->_gating_mode == GatingMode::NONE)
   {
     this->_activation->apply(this->_z.leftCols(num_frames));
     _1x1.process_(_z, num_frames);
   }
-  else
+  else if (this->_gating_mode == GatingMode::GATED)
   {
-    // CAREFUL: .topRows() and .bottomRows() won't be memory-contiguous for a column-major matrix (Issue 125). Need to
-    // do this column-wise:
-    for (int i = 0; i < num_frames; i++)
-    {
-      this->_activation->apply(this->_z.block(0, i, bottleneck, 1));
-      // TODO Need to support other activation functions here instead of hardcoded sigmoid
-      activations::Activation::get_activation("Sigmoid")->apply(this->_z.block(bottleneck, i, bottleneck, 1));
-    }
-    this->_z.block(0, 0, bottleneck, num_frames).array() *=
-      this->_z.block(bottleneck, 0, bottleneck, num_frames).array();
-    _1x1.process_(_z.topRows(bottleneck), num_frames); // Might not be RT safe
+    // Use the GatingActivation class
+    // Extract the blocks first to avoid temporary reference issues
+    auto input_block = this->_z.leftCols(num_frames);
+    auto output_block = this->_z.topRows(bottleneck).leftCols(num_frames);
+    this->_gating_activation->apply(input_block, output_block);
+    _1x1.process_(this->_z.topRows(bottleneck), num_frames);
+  }
+  else if (this->_gating_mode == GatingMode::BLENDED)
+  {
+    // Use the BlendingActivation class
+    // Extract the blocks first to avoid temporary reference issues
+    auto input_block = this->_z.leftCols(num_frames);
+    auto output_block = this->_z.topRows(bottleneck).leftCols(num_frames);
+    this->_blending_activation->apply(input_block, output_block);
+    _1x1.process_(this->_z.topRows(bottleneck), num_frames);
   }
 
   if (this->_head1x1)
   {
-    if (!this->_gated)
+    if (this->_gating_mode == GatingMode::NONE)
       this->_head1x1->process_(this->_z.leftCols(num_frames), num_frames);
     else
       this->_head1x1->process(this->_z.topRows(bottleneck).leftCols(num_frames), num_frames);
@@ -89,7 +98,7 @@ void nam::wavenet::_Layer::Process(const Eigen::MatrixXf& input, const Eigen::Ma
   else
   {
     // Store output to head (skip connection: activated conv output)
-    if (!this->_gated)
+    if (this->_gating_mode == GatingMode::NONE)
       this->_output_head.leftCols(num_frames).noalias() = this->_z.leftCols(num_frames);
     else
       this->_output_head.leftCols(num_frames).noalias() = this->_z.topRows(bottleneck).leftCols(num_frames);
@@ -105,15 +114,16 @@ void nam::wavenet::_Layer::Process(const Eigen::MatrixXf& input, const Eigen::Ma
 nam::wavenet::_LayerArray::_LayerArray(const int input_size, const int condition_size, const int head_size,
                                        const int channels, const int bottleneck, const int kernel_size,
                                        const std::vector<int>& dilations, const std::string activation,
-                                       const bool gated, const bool head_bias, const int groups_input,
-                                       const int groups_1x1, const Head1x1Params& head1x1_params)
+                                       const GatingMode gating_mode, const bool head_bias, const int groups_input,
+                                       const int groups_1x1, const Head1x1Params& head1x1_params,
+                                       const std::string& secondary_activation)
 : _rechannel(input_size, channels, false)
 , _head_rechannel(bottleneck, head_size, head_bias)
 , _bottleneck(bottleneck)
 {
   for (size_t i = 0; i < dilations.size(); i++)
-    this->_layers.push_back(_Layer(condition_size, channels, bottleneck, kernel_size, dilations[i], activation, gated,
-                                   groups_input, groups_1x1, head1x1_params));
+    this->_layers.push_back(_Layer(condition_size, channels, bottleneck, kernel_size, dilations[i], activation,
+                                   gating_mode, groups_input, groups_1x1, head1x1_params, secondary_activation));
 }
 
 void nam::wavenet::_LayerArray::SetMaxBufferSize(const int maxBufferSize)
@@ -263,9 +273,9 @@ nam::wavenet::WaveNet::WaveNet(const int in_channels,
     this->_layer_arrays.push_back(nam::wavenet::_LayerArray(
       layer_array_params[i].input_size, layer_array_params[i].condition_size, layer_array_params[i].head_size,
       layer_array_params[i].channels, layer_array_params[i].bottleneck, layer_array_params[i].kernel_size,
-      layer_array_params[i].dilations, layer_array_params[i].activation, layer_array_params[i].gated,
+      layer_array_params[i].dilations, layer_array_params[i].activation, layer_array_params[i].gating_mode,
       layer_array_params[i].head_bias, layer_array_params[i].groups_input, layer_array_params[i].groups_1x1,
-      layer_array_params[i].head1x1_params));
+      layer_array_params[i].head1x1_params, layer_array_params[i].secondary_activation));
     if (i > 0)
       if (layer_array_params[i].channels != layer_array_params[i - 1].head_size)
       {
@@ -468,7 +478,50 @@ std::unique_ptr<nam::DSP> nam::wavenet::Factory(const nlohmann::json& config, st
     const int kernel_size = layer_config["kernel_size"];
     const auto dilations = layer_config["dilations"];
     const std::string activation = layer_config["activation"].get<std::string>();
-    const bool gated = layer_config["gated"];
+    // Parse gating mode - support both old "gated" boolean and new "gating_mode" string
+    GatingMode gating_mode = GatingMode::NONE;
+    std::string secondary_activation;
+
+    if (layer_config.find("gating_mode") != layer_config.end())
+    {
+      std::string gating_mode_str = layer_config["gating_mode"].get<std::string>();
+      if (gating_mode_str == "gated")
+      {
+        gating_mode = GatingMode::GATED;
+        secondary_activation = layer_config["secondary_activation"].get<std::string>();
+      }
+      else if (gating_mode_str == "blended")
+      {
+        gating_mode = GatingMode::BLENDED;
+        secondary_activation = layer_config["secondary_activation"].get<std::string>();
+      }
+      else if (gating_mode_str == "none")
+      {
+        gating_mode = GatingMode::NONE;
+        secondary_activation.clear();
+      }
+      else
+        throw std::runtime_error("Invalid gating_mode: " + gating_mode_str);
+    }
+    else if (layer_config.find("gated") != layer_config.end())
+    {
+      // Backward compatibility: convert old "gated" boolean to new enum
+      bool gated = layer_config["gated"];
+      gating_mode = gated ? GatingMode::GATED : GatingMode::NONE;
+      if (gated)
+      {
+        secondary_activation = "Sigmoid";
+      }
+      else
+      {
+        secondary_activation.clear();
+      }
+    }
+    else
+    {
+      throw std::invalid_argument("No information on gating mode found for layer array " + std::to_string(i));
+    }
+
     const bool head_bias = layer_config["head_bias"];
 
     // Parse head1x1 parameters
@@ -477,9 +530,9 @@ std::unique_ptr<nam::DSP> nam::wavenet::Factory(const nlohmann::json& config, st
     int head1x1_groups = layer_config.value("head1x1_groups", 1);
     nam::wavenet::Head1x1Params head1x1_params(head1x1_active, head1x1_out_channels, head1x1_groups);
 
-    layer_array_params.push_back(nam::wavenet::LayerArrayParams(input_size, condition_size, head_size, channels,
-                                                                bottleneck, kernel_size, dilations, activation, gated,
-                                                                head_bias, groups, groups_1x1, head1x1_params));
+    layer_array_params.push_back(nam::wavenet::LayerArrayParams(
+      input_size, condition_size, head_size, channels, bottleneck, kernel_size, dilations, activation, gating_mode,
+      head_bias, groups, groups_1x1, head1x1_params, secondary_activation));
   }
   const bool with_head = !config["head"].is_null();
   const float head_scale = config["head_scale"];

NAM/wavenet.h

@@ -3,17 +3,33 @@
 #include <string>
 #include <vector>
 #include <memory>
+#include <stdexcept>
 
 #include "json.hpp"
 #include <Eigen/Dense>
 
 #include "dsp.h"
 #include "conv1d.h"
+#include "gating_activations.h"
 
 namespace nam
 {
 namespace wavenet
 {
+
+// Gating mode for WaveNet layers
+enum class GatingMode
+{
+  NONE, // No gating or blending
+  GATED, // Traditional gating (element-wise multiplication)
+  BLENDED // Blending (weighted average)
+};
+
+// Helper function for backward compatibility with boolean gated parameter
+inline GatingMode gating_mode_from_bool(bool gated)
+{
+  return gated ? GatingMode::GATED : GatingMode::NONE;
+}
 // Parameters for head1x1 configuration
 struct Head1x1Params
 {
@@ -32,20 +48,42 @@ struct Head1x1Params
 class _Layer
 {
 public:
+  // New constructor with GatingMode enum and configurable activations
   _Layer(const int condition_size, const int channels, const int bottleneck, const int kernel_size, const int dilation,
-         const std::string activation, const bool gated, const int groups_input, const int groups_1x1,
-         const Head1x1Params& head1x1_params)
-  : _conv(channels, gated ? 2 * bottleneck : bottleneck, kernel_size, true, dilation)
-  , _input_mixin(condition_size, gated ? 2 * bottleneck : bottleneck, false)
+         const std::string activation, const GatingMode gating_mode, const int groups_input, const int groups_1x1,
+         const Head1x1Params& head1x1_params, const std::string& secondary_activation)
+  : _conv(channels, (gating_mode != GatingMode::NONE) ? 2 * bottleneck : bottleneck, kernel_size, true, dilation)
+  , _input_mixin(condition_size, (gating_mode != GatingMode::NONE) ? 2 * bottleneck : bottleneck, false)
   , _1x1(bottleneck, channels, groups_1x1)
   , _activation(activations::Activation::get_activation(activation)) // needs to support activations with parameters
-  , _gated(gated)
+  , _gating_mode(gating_mode)
   , _bottleneck(bottleneck)
   {
     if (head1x1_params.active)
     {
       _head1x1 = std::make_unique<Conv1x1>(bottleneck, head1x1_params.out_channels, true, head1x1_params.groups);
     }
+
+    // Validate & initialize gating/blending activation
+    if (gating_mode == GatingMode::GATED)
+    {
+      if (secondary_activation.empty())
+        throw std::invalid_argument("secondary_activation must be provided for gated mode");
+      _gating_activation = std::make_unique<gating_activations::GatingActivation>(
+        _activation, activations::Activation::get_activation(secondary_activation), bottleneck);
+    }
+    else if (gating_mode == GatingMode::BLENDED)
+    {
+      if (secondary_activation.empty())
+        throw std::invalid_argument("secondary_activation must be provided for blended mode");
+      _blending_activation = std::make_unique<gating_activations::BlendingActivation>(
+        _activation, activations::Activation::get_activation(secondary_activation), bottleneck);
+    }
+    else
+    {
+      if (!secondary_activation.empty())
+        throw std::invalid_argument("secondary_activation provided for none mode");
+    }
   };
 
   // Resize all arrays to be able to process `maxBufferSize` frames.
@@ -97,17 +135,22 @@ class _Layer
   Eigen::MatrixXf _output_head;
 
   activations::Activation* _activation;
-  const bool _gated;
+  const GatingMode _gating_mode;
   const int _bottleneck; // Internal channel count (not doubled when gated)
+
+  // Gating/blending activation objects
+  std::unique_ptr<gating_activations::GatingActivation> _gating_activation;
+  std::unique_ptr<gating_activations::BlendingActivation> _blending_activation;
 };
 
 class LayerArrayParams
 {
 public:
   LayerArrayParams(const int input_size_, const int condition_size_, const int head_size_, const int channels_,
                    const int bottleneck_, const int kernel_size_, const std::vector<int>&& dilations_,
-                   const std::string activation_, const bool gated_, const bool head_bias_, const int groups_input,
-                   const int groups_1x1_, const Head1x1Params& head1x1_params_)
+                   const std::string activation_, const GatingMode gating_mode_, const bool head_bias_,
+                   const int groups_input, const int groups_1x1_, const Head1x1Params& head1x1_params_,
+                   const std::string& secondary_activation_)
   : input_size(input_size_)
   , condition_size(condition_size_)
   , head_size(head_size_)
@@ -116,11 +159,12 @@ class LayerArrayParams
   , kernel_size(kernel_size_)
   , dilations(std::move(dilations_))
   , activation(activation_)
-  , gated(gated_)
+  , gating_mode(gating_mode_)
   , head_bias(head_bias_)
   , groups_input(groups_input)
   , groups_1x1(groups_1x1_)
   , head1x1_params(head1x1_params_)
+  , secondary_activation(secondary_activation_)
   {
   }
 
@@ -132,21 +176,23 @@ class LayerArrayParams
   const int kernel_size;
   std::vector<int> dilations;
   const std::string activation;
-  const bool gated;
+  const GatingMode gating_mode;
   const bool head_bias;
   const int groups_input;
   const int groups_1x1;
   const Head1x1Params head1x1_params;
+  const std::string secondary_activation;
 };
 
 // An array of layers with the same channels, kernel sizes, activations.
 class _LayerArray
 {
 public:
+  // New constructor with GatingMode enum and configurable activations
   _LayerArray(const int input_size, const int condition_size, const int head_size, const int channels,
               const int bottleneck, const int kernel_size, const std::vector<int>& dilations,
-              const std::string activation, const bool gated, const bool head_bias, const int groups_input,
-              const int groups_1x1, const Head1x1Params& head1x1_params);
+              const std::string activation, const GatingMode gating_mode, const bool head_bias, const int groups_input,
+              const int groups_1x1, const Head1x1Params& head1x1_params, const std::string& secondary_activation);
 
   void SetMaxBufferSize(const int maxBufferSize);
 

tools/run_tests.cpp

@@ -21,6 +21,7 @@
 #include "test/test_blending_detailed.cpp"
 #include "test/test_input_buffer_verification.cpp"
 #include "test/test_lstm.cpp"
+#include "test/test_wavenet_configurable_gating.cpp"
 
 int main()
 {
@@ -177,6 +178,9 @@ int main()
   test_input_buffer_verification::TestInputBufferVerification::test_buffer_stores_pre_activation_values();
   test_input_buffer_verification::TestInputBufferVerification::test_buffer_with_different_activations();
 
+  // Configurable gating/blending tests
+  run_configurable_gating_tests();
+
   test_get_dsp::test_gets_input_level();
   test_get_dsp::test_gets_output_level();
   test_get_dsp::test_null_input_level();