diff --git a/dlib/cuda/cpu_dlib.cpp b/dlib/cuda/cpu_dlib.cpp index b7cfbde026..0de9ae9b31 100644 --- a/dlib/cuda/cpu_dlib.cpp +++ b/dlib/cuda/cpu_dlib.cpp @@ -3219,6 +3219,153 @@ namespace dlib // ------------------------------------------------------------------------------------ + void compute_act_halt_probabilities( + resizable_tensor& halt_probs, + resizable_tensor& logits, + const tensor& input_data, + const tensor& halt_params, + long batch_size, + long seq_len, + long feature_dim + ) + { + const float* in_ptr = input_data.host(); + const float* W_halt = halt_params.host(); + const float b_halt = halt_params.host()[feature_dim]; + float* logits_ptr = logits.host(); + float* halt_probs_ptr = halt_probs.host(); + + const long d_model = feature_dim / input_data.k(); + const long num_channels = input_data.k(); + + for (long pos = 0; pos < batch_size * seq_len; ++pos) { + const long n = pos / seq_len; + const long s = pos % seq_len; + + float logit = b_halt; + + for (long c = 0; c < num_channels; ++c) { + for (long d = 0; d < d_model; ++d) { + const long in_idx = ((n * num_channels + c) * seq_len + s) * d_model + d; + const long weight_idx = c * d_model + d; + logit += in_ptr[in_idx] * W_halt[weight_idx]; + } + } + + logits_ptr[pos] = logit; + + halt_probs_ptr[pos] = 1.0f / (1.0f + std::exp(-logit)); + } + } + + void update_act_state( + resizable_tensor& output, + const tensor& input_data, + const tensor& halt_probs, + resizable_tensor& cumulative_halting, + resizable_tensor& remainders, + resizable_tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float halt_threshold, + long current_step + ) + { + const float* in_ptr = input_data.host(); + const float* p_halt = halt_probs.host(); + float* out_ptr = output.host(); + float* cum_halt = cumulative_halting.host(); + float* remain = remainders.host(); + float* steps = n_steps.host(); + + for (long pos = 0; pos < batch_size * seq_len; ++pos) { + if (cum_halt[pos] < halt_threshold) { + const long n = pos / seq_len; + const long s = pos % seq_len; + + float p = p_halt[pos]; + float r = remain[pos]; + float effective = std::min(p * r, halt_threshold - cum_halt[pos]); + + cum_halt[pos] += effective; + remain[pos] -= effective; + steps[pos] = static_cast(current_step + 1); + + for (long c = 0; c < num_channels; ++c) { + for (long d = 0; d < d_model; ++d) { + const long idx = ((n * num_channels + c) * seq_len + s) * d_model + d; + out_ptr[idx] += effective * in_ptr[idx]; + } + } + } + } + } + + void finalize_act_output( + resizable_tensor& output, + const tensor& input_data, + const tensor& remainders, + long batch_size, + long seq_len, + long d_model, + long num_channels + ) + { + const float* in_ptr = input_data.host(); + const float* remain = remainders.host(); + float* out_ptr = output.host(); + + for (long pos = 0; pos < batch_size * seq_len; ++pos) { + float r = remain[pos]; + if (r > 1e-6f) { + const long n = pos / seq_len; + const long s = pos % seq_len; + + for (long c = 0; c < num_channels; ++c) { + for (long d = 0; d < d_model; ++d) { + const long idx = ((n * num_channels + c) * seq_len + s) * d_model + d; + out_ptr[idx] += r * in_ptr[idx]; + } + } + } + } + } + + void apply_act_depth_scaling( + tensor& gradients, + const tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float max_steps, + float scale_factor + ) + { + const float* steps = n_steps.host(); + float* grad_ptr = gradients.host(); + + for (long pos = 0; pos < batch_size * seq_len; ++pos) + { + const float scale = 1.0f + scale_factor * (steps[pos] / max_steps); + const long n = pos / seq_len; + const long s = pos % seq_len; + + for (long c = 0; c < num_channels; ++c) + { + for (long d = 0; d < d_model; ++d) + { + const long idx = ((n * num_channels + c) * seq_len + s) * d_model + d; + grad_ptr[idx] *= scale; + } + } + } + } + + // ------------------------------------------------------------------------------------ + } } diff --git a/dlib/cuda/cpu_dlib.h b/dlib/cuda/cpu_dlib.h index ab88a4a4ee..4e29c8a8d9 100644 --- a/dlib/cuda/cpu_dlib.h +++ b/dlib/cuda/cpu_dlib.h @@ -536,6 +536,54 @@ namespace dlib bool scale ); + // ----------------------------------------------------------------------------------- + + void compute_act_halt_probabilities( + resizable_tensor& halt_probs, + resizable_tensor& logits, + const tensor& input_data, + const tensor& halt_params, + long batch_size, + long seq_len, + long feature_dim + ); + + void update_act_state( + resizable_tensor& output, + const tensor& input_data, + const tensor& halt_probs, + resizable_tensor& cumulative_halting, + resizable_tensor& remainders, + resizable_tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float halt_threshold, + long current_step + ); + + void finalize_act_output( + resizable_tensor& output, + const tensor& input_data, + const tensor& remainders, + long batch_size, + long seq_len, + long d_model, + long num_channels + ); + + void apply_act_depth_scaling( + tensor& gradients, + const tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float max_steps, + float scale_factor + ); + // ----------------------------------------------------------------------------------- class pooling diff --git a/dlib/cuda/cuda_dlib.cu b/dlib/cuda/cuda_dlib.cu index 5d6ec4052c..672efe9c22 100644 --- a/dlib/cuda/cuda_dlib.cu +++ b/dlib/cuda/cuda_dlib.cu @@ -1,4 +1,4 @@ -// Copyright (C) 2015 Davis E. King (davis@dlib.net) +// Copyright (C) 2015 Davis E. King (davis@dlib.net) // License: Boost Software License See LICENSE.txt for the full license. #include "cuda_utils.h" @@ -2733,6 +2733,264 @@ namespace dlib src.k(), src.nr(), src.nc(), src.device(), add_to); } + // ---------------------------------------------------------------------------------------- + + // CUDA Kernels for ACT operations + __global__ void _cuda_compute_act_halt_probabilities( + float* halt_probs, + float* logits, + const float* input_data, + const float* W_halt, + float b_halt, + size_t batch_size, + size_t seq_len, + size_t d_model, + size_t num_channels, + size_t feature_dim + ) + { + const long total_positions = batch_size * seq_len; + + for (auto pos : grid_stride_range_y(0, total_positions)) + for (auto i : grid_stride_range(0, 1)) + logits[pos] = b_halt; + __syncthreads(); + + for (auto pos : grid_stride_range_y(0, total_positions)) + { + const long n = pos / seq_len; + const long s = pos % seq_len; + + float temp = 0; + for (auto feat_idx : grid_stride_range(0, feature_dim)) + { + const long c = feat_idx / d_model; + const long d = feat_idx % d_model; + + const long in_idx = ((n * num_channels + c) * seq_len + s) * d_model + d; + temp += input_data[in_idx] * W_halt[feat_idx]; + } + + warp_reduce_atomic_add(logits[pos], temp); + } + __syncthreads(); + + for (auto pos : grid_stride_range(0, total_positions)) + { + halt_probs[pos] = 1.0f / (1.0f + expf(-logits[pos])); + } + } + + void compute_act_halt_probabilities( + resizable_tensor& halt_probs, + resizable_tensor& logits, + const tensor& input_data, + const tensor& halt_params, + long batch_size, + long seq_len, + long feature_dim + ) + { + const long total_positions = batch_size * seq_len; + const long d_model = feature_dim / input_data.k(); + const long num_channels = input_data.k(); + + halt_probs.set_size(total_positions, 1, 1, 1); + logits.set_size(total_positions, 1, 1, 1); + + launch_kernel(_cuda_compute_act_halt_probabilities, + max_jobs(feature_dim, total_positions), + halt_probs.device(), + logits.device(), + input_data.device(), + halt_params.device(), + halt_params.host()[feature_dim], + batch_size, + seq_len, + d_model, + num_channels, + feature_dim); + } + + __global__ void _cuda_update_act_state( + float* output, + const float* input_data, + const float* halt_probs, + float* cumulative_halting, + float* remainders, + float* n_steps, + size_t batch_size, + size_t seq_len, + size_t d_model, + size_t num_channels, + float halt_threshold, + long current_step + ) + { + for (auto pos : grid_stride_range(0, batch_size * seq_len)) + { + if (cumulative_halting[pos] < halt_threshold) + { + const size_t n = pos / seq_len; + const size_t s = pos % seq_len; + + float p = halt_probs[pos]; + float r = remainders[pos]; + float effective = fminf(p * r, halt_threshold - cumulative_halting[pos]); + + cumulative_halting[pos] += effective; + remainders[pos] -= effective; + n_steps[pos] = static_cast(current_step + 1); + + for (size_t c = 0; c < num_channels; ++c) { + for (size_t d = 0; d < d_model; ++d) { + const size_t idx = ((n * num_channels + c) * seq_len + s) * d_model + d; + output[idx] += effective * input_data[idx]; + } + } + } + } + } + + void update_act_state( + resizable_tensor& output, + const tensor& input_data, + const tensor& halt_probs, + resizable_tensor& cumulative_halting, + resizable_tensor& remainders, + resizable_tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float halt_threshold, + long current_step + ) + { + const long total_positions = batch_size * seq_len; + + launch_kernel(_cuda_update_act_state, + max_jobs(total_positions), + output.device(), + input_data.device(), + halt_probs.device(), + cumulative_halting.device(), + remainders.device(), + n_steps.device(), + batch_size, + seq_len, + d_model, + num_channels, + halt_threshold, + current_step); + } + + __global__ void _cuda_finalize_act_output( + float* output, + const float* input_data, + const float* remainders, + size_t batch_size, + size_t seq_len, + size_t d_model, + size_t num_channels + ) + { + for (auto pos : grid_stride_range(0, batch_size * seq_len)) + { + float r = remainders[pos]; + if (r > 1e-6f) { + const size_t n = pos / seq_len; + const size_t s = pos % seq_len; + + for (size_t c = 0; c < num_channels; ++c) { + for (size_t d = 0; d < d_model; ++d) { + const size_t idx = ((n * num_channels + c) * seq_len + s) * d_model + d; + output[idx] += r * input_data[idx]; + } + } + } + } + } + + void finalize_act_output( + resizable_tensor& output, + const tensor& input_data, + const tensor& remainders, + long batch_size, + long seq_len, + long d_model, + long num_channels + ) + { + const long total_positions = batch_size * seq_len; + + launch_kernel(_cuda_finalize_act_output, + max_jobs(total_positions), + output.device(), + input_data.device(), + remainders.device(), + batch_size, + seq_len, + d_model, + num_channels); + } + + __global__ void _cuda_apply_act_depth_scaling( + float* gradients, + const float* n_steps, + size_t batch_size, + size_t seq_len, + size_t d_model, + size_t num_channels, + float max_steps, + float scale_factor + ) + { + const long total_positions = batch_size * seq_len; + const long feature_dim = num_channels * d_model; + + for (auto pos : grid_stride_range_y(0, total_positions)) + { + const long n = pos / seq_len; + const long s = pos % seq_len; + const float scale = 1.0f + scale_factor * (n_steps[pos] / max_steps); + + for (auto feat_idx : grid_stride_range(0, feature_dim)) + { + const long c = feat_idx / d_model; + const long d = feat_idx % d_model; + const long idx = ((n * num_channels + c) * seq_len + s) * d_model + d; + gradients[idx] *= scale; + } + } + } + + void apply_act_depth_scaling( + tensor& gradients, + const tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float max_steps, + float scale_factor + ) + { + const long total_positions = batch_size * seq_len; + const long feature_dim = num_channels * d_model; + + launch_kernel(_cuda_apply_act_depth_scaling, + max_jobs(feature_dim, total_positions), + gradients.device(), + n_steps.device(), + batch_size, + seq_len, + d_model, + num_channels, + max_steps, + scale_factor); + } + // ---------------------------------------------------------------------------------------- diff --git a/dlib/cuda/cuda_dlib.h b/dlib/cuda/cuda_dlib.h index 2f22b7e23e..26e1d29e4f 100644 --- a/dlib/cuda/cuda_dlib.h +++ b/dlib/cuda/cuda_dlib.h @@ -608,6 +608,54 @@ namespace dlib const tensor& src ); + // ---------------------------------------------------------------------------------------- + + void compute_act_halt_probabilities( + resizable_tensor& halt_probs, + resizable_tensor& logits, + const tensor& input_data, + const tensor& halt_params, + long batch_size, + long seq_len, + long feature_dim + ); + + void update_act_state( + resizable_tensor& output, + const tensor& input_data, + const tensor& halt_probs, + resizable_tensor& cumulative_halting, + resizable_tensor& remainders, + resizable_tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float halt_threshold, + long current_step + ); + + void finalize_act_output( + resizable_tensor& output, + const tensor& input_data, + const tensor& remainders, + long batch_size, + long seq_len, + long d_model, + long num_channels + ); + + void apply_act_depth_scaling( + tensor& gradients, + const tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float max_steps, + float scale_factor + ); + // ---------------------------------------------------------------------------------------- class compute_loss_binary_log_per_pixel diff --git a/dlib/cuda/tensor_tools.cpp b/dlib/cuda/tensor_tools.cpp index 8dece8369f..d9429df2f4 100644 --- a/dlib/cuda/tensor_tools.cpp +++ b/dlib/cuda/tensor_tools.cpp @@ -1412,6 +1412,90 @@ namespace dlib { namespace tt #endif } +// ---------------------------------------------------------------------------------------- + + void compute_act_halt_probabilities( + resizable_tensor& halt_probs, + resizable_tensor& logits, + const tensor& input_data, + const tensor& halt_params, + long batch_size, + long seq_len, + long feature_dim + ) + { +#ifdef DLIB_USE_CUDA + cuda::compute_act_halt_probabilities(halt_probs, logits, input_data, halt_params, + batch_size, seq_len, feature_dim); +#else + cpu::compute_act_halt_probabilities(halt_probs, logits, input_data, halt_params, + batch_size, seq_len, feature_dim); +#endif + } + + void update_act_state( + resizable_tensor& output, + const tensor& input_data, + const tensor& halt_probs, + resizable_tensor& cumulative_halting, + resizable_tensor& remainders, + resizable_tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float halt_threshold, + long current_step + ) + { +#ifdef DLIB_USE_CUDA + cuda::update_act_state(output, input_data, halt_probs, cumulative_halting, remainders, + n_steps, batch_size, seq_len, d_model, num_channels, halt_threshold, current_step); +#else + cpu::update_act_state(output, input_data, halt_probs, cumulative_halting, remainders, + n_steps, batch_size, seq_len, d_model, num_channels, halt_threshold, current_step); +#endif + } + + void finalize_act_output( + resizable_tensor& output, + const tensor& input_data, + const tensor& remainders, + long batch_size, + long seq_len, + long d_model, + long num_channels + ) + { +#ifdef DLIB_USE_CUDA + cuda::finalize_act_output(output, input_data, remainders, + batch_size, seq_len, d_model, num_channels); +#else + cpu::finalize_act_output(output, input_data, remainders, + batch_size, seq_len, d_model, num_channels); +#endif + } + + void apply_act_depth_scaling( + tensor& gradients, + const tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float max_steps, + float scale_factor + ) + { +#ifdef DLIB_USE_CUDA + cuda::apply_act_depth_scaling(gradients, n_steps, batch_size, seq_len, + d_model, num_channels, max_steps, scale_factor); +#else + cpu::apply_act_depth_scaling(gradients, n_steps, batch_size, seq_len, + d_model, num_channels, max_steps, scale_factor); +#endif + } + // ---------------------------------------------------------------------------------------- }} diff --git a/dlib/cuda/tensor_tools.h b/dlib/cuda/tensor_tools.h index 18a5564f98..fe0260ea88 100644 --- a/dlib/cuda/tensor_tools.h +++ b/dlib/cuda/tensor_tools.h @@ -1,4 +1,4 @@ -// Copyright (C) 2015 Davis E. King (davis@dlib.net) +// Copyright (C) 2015 Davis E. King (davis@dlib.net) // License: Boost Software License See LICENSE.txt for the full license. #ifndef DLIB_TeNSOR_TOOLS_H_ #define DLIB_TeNSOR_TOOLS_H_ @@ -2392,6 +2392,130 @@ namespace dlib { namespace tt - #dest(n,k,c,r) == dest(n,k,c,r) + src(n,k,r,c) !*/ +// ---------------------------------------------------------------------------------------- + + // ACT (Adaptive Computation Time) operations + + void compute_act_halt_probabilities( + resizable_tensor& halt_probs, + resizable_tensor& logits, + const tensor& input_data, + const tensor& halt_params, + long batch_size, + long seq_len, + long feature_dim + ); + /*! + requires + - halt_params.size() == feature_dim + 1 (weights + bias) + - input_data.num_samples() == batch_size + - input_data.k() == num_channels where feature_dim = num_channels * d_model + - input_data.nr() == seq_len + - input_data.nc() == d_model + ensures + - Computes halting probabilities for Adaptive Computation Time: + - halt_probs contains sigmoid(W_halt^T * input + b_halt) for each position + - logits contains the pre-sigmoid values + - batch_size: number of samples in the batch + - seq_len: sequence length (number of positions to process) + - feature_dim: total feature dimension (num_channels × d_model) + !*/ + + void update_act_state( + resizable_tensor& output, + const tensor& input_data, + const tensor& halt_probs, + resizable_tensor& cumulative_halting, + resizable_tensor& remainders, + resizable_tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float halt_threshold, + long current_step + ); + /*! + requires + - 0 < halt_threshold <= 1.0 + - current_step >= 0 + - input_data.num_samples() == batch_size + - input_data.k() == num_channels + - input_data.nr() == seq_len + - input_data.nc() == d_model + - output has the same dimensions as input_data + - halt_probs.size() == batch_size * seq_len + - cumulative_halting.size() == remainders.size() == n_steps.size() == batch_size * seq_len + ensures + - Core ACT update step that accumulates weighted outputs: + - Updates ACT state for all positions + - Accumulates weighted outputs: output += α_t^n * input_data + - Updates cumulative_halting, remainders, and n_steps + - batch_size: number of samples in the batch + - seq_len: sequence length (number of positions to process) + - d_model: model dimension per channel + - num_channels: number of feature channels + - halt_threshold: halting threshold (typically 0.99) + - current_step: current computation step index (0-based) + !*/ + + void finalize_act_output( + resizable_tensor& output, + const tensor& input_data, + const tensor& remainders, + long batch_size, + long seq_len, + long d_model, + long num_channels + ); + /*! + requires + - input_data.num_samples() == batch_size + - input_data.k() == num_channels + - input_data.nr() == seq_len + - input_data.nc() == d_model + - output has the same dimensions as input_data + - remainders.size() == batch_size * seq_len + ensures + - Finalizes ACT output by adding remainder contributions: + - Adds final remainder contributions: output += ρ_t * input_data + - Applied only to positions with significant remainder (> 1e-6) + - batch_size: number of samples in the batch + - seq_len: sequence length (number of positions to process) + - d_model: model dimension per channel + - num_channels: number of feature channels + !*/ + + void apply_act_depth_scaling( + tensor& gradients, + const tensor& n_steps, + long batch_size, + long seq_len, + long d_model, + long num_channels, + float max_steps, + float scale_factor + ); + /*! + requires + - scale_factor >= 0 + - max_steps > 0 + - gradients.num_samples() == batch_size + - gradients.k() == num_channels + - gradients.nr() == seq_len + - gradients.nc() == d_model + - n_steps.size() == batch_size * seq_len + ensures + - Applies gradient scaling based on computation depth: + - Applies depth-dependent gradient scaling + - scale = 1 + scale_factor * (n_steps[pos] / max_steps) + - seq_len: sequence length (number of positions to process) + - d_model: model dimension per channel + - num_channels: number of feature channels + - max_steps: maximum allowed computation steps + - scale_factor: scaling strength (0 = no scaling) + !*/ + // ---------------------------------------------------------------------------------------- }} diff --git a/dlib/dnn/layers.h b/dlib/dnn/layers.h index 0ff4d2c301..6f9389fced 100644 --- a/dlib/dnn/layers.h +++ b/dlib/dnn/layers.h @@ -1,4 +1,4 @@ -// Copyright (C) 2015 Davis E. King (davis@dlib.net) +// Copyright (C) 2015 Davis E. King (davis@dlib.net) // License: Boost Software License See LICENSE.txt for the full license. #ifndef DLIB_DNn_LAYERS_H_ #define DLIB_DNn_LAYERS_H_ @@ -5711,6 +5711,402 @@ namespace dlib template using tril_diag = add_layer, SUBNET>; +// ---------------------------------------------------------------------------------------- + + template + class adaptive_computation_time_ { + public: + explicit adaptive_computation_time_() : + max_steps_(max_steps), + halt_threshold_(0.99f), // theta in Graves' notation + ponder_penalty_(0.01f), // lambda (ponder cost weight) + enable_depth_scaling_(false), + batch_size_(0), + seq_len_(0), + d_model_(0), + num_channels_(0), + feature_dim_(0), + ponder_cost_(0), + avg_steps_(0) + { + } + + adaptive_computation_time_(const adaptive_computation_time_& item) : + max_steps_(item.max_steps_), + halt_threshold_(item.halt_threshold_), + ponder_penalty_(item.ponder_penalty_), + enable_depth_scaling_(item.enable_depth_scaling_), + batch_size_(item.batch_size_), + seq_len_(item.seq_len_), + d_model_(item.d_model_), + num_channels_(item.num_channels_), + feature_dim_(item.feature_dim_), + ponder_cost_(item.ponder_cost_), + avg_steps_(item.avg_steps_), + params(item.params), + halting_probs_(item.halting_probs_), + cumulative_halting_(item.cumulative_halting_), + remainders_(item.remainders_), + n_steps_(item.n_steps_), + logits_(item.logits_), + grad_logits_(item.grad_logits_), + input_cache_(item.input_cache_), + true_effective_weights_(item.true_effective_weights_) + { + } + + adaptive_computation_time_& operator=(const adaptive_computation_time_& item) + { + if (this == &item) + return *this; + + max_steps_ = item.max_steps_; + halt_threshold_ = item.halt_threshold_; + ponder_penalty_ = item.ponder_penalty_; + enable_depth_scaling_ = item.enable_depth_scaling_; + batch_size_ = item.batch_size_; + seq_len_ = item.seq_len_; + d_model_ = item.d_model_; + num_channels_ = item.num_channels_; + feature_dim_ = item.feature_dim_; + ponder_cost_ = item.ponder_cost_; + avg_steps_ = item.avg_steps_; + params = item.params; + halting_probs_ = item.halting_probs_; + cumulative_halting_ = item.cumulative_halting_; + remainders_ = item.remainders_; + n_steps_ = item.n_steps_; + logits_ = item.logits_; + grad_logits_ = item.grad_logits_; + input_cache_ = item.input_cache_; + true_effective_weights_ = item.true_effective_weights_; + + return *this; + } + + template + void setup(const SUBNET& sub) { + const auto& input = sub.get_output(); + + // Store expected dimensions for parameter initialization + batch_size_ = input.num_samples(); + seq_len_ = input.nr(); + d_model_ = input.nc(); + num_channels_ = input.k(); + feature_dim_ = d_model_ * num_channels_; + + // Initialize halting parameters + params.set_size(1, 1, feature_dim_ + 1, 1); + + // He initialization for stability + dlib::rand rnd(std::rand()); + const float scale = std::sqrt(2.0f / feature_dim_); + float* p = params.host(); + + // Initialize weight matrix W_halt + for (long i = 0; i < feature_dim_; ++i) + p[i] = rnd.get_random_gaussian() * scale; + + // Initialize bias b_halt (typically zero) + p[feature_dim_] = 0.0f; + + // Pre-allocate workspace for maximum expected size + allocate_workspace(); + } + + template + void forward(const SUBNET& sub, resizable_tensor& output) { + const tensor& input = sub.get_output(); + output.copy_size(input); + + // Ensure workspace is allocated for current batch dimensions + const long curr_batch = input.num_samples(); + if (curr_batch != batch_size_) { + batch_size_ = curr_batch; + allocate_workspace(); + } + + // Initialize output for weighted accumulation + output = 0; + + // Initialize ACT state vectors + const long total_positions = batch_size_ * seq_len_; + float* cum_halt_ptr = cumulative_halting_.host(); + float* remainders_ptr = remainders_.host(); + float* n_steps_ptr = n_steps_.host(); + + for (long i = 0; i < total_positions; ++i) { + cum_halt_ptr[i] = 0.0f; // h_t^n: cumulative halting probability + remainders_ptr[i] = 1.0f; // ρ_t: remaining probability mass + n_steps_ptr[i] = 0.0f; // N(t): number of computation steps + } + + // Cache input for backward pass + input_cache_.copy_size(input); + tt::copy_tensor(false, input_cache_, 0, input, 0, input.k()); + + // Initialize effective weights tracker for gradient computation + true_effective_weights_.set_size(total_positions, 1, 1, 1); + true_effective_weights_ = 0; + + // Main ACT computation loop + for (long step = 0; step < max_steps_; ++step) { + + // Compute halting probabilities: p_t^n = sigmoid(W_halt^T * s_t^n + b_halt) + tt::compute_act_halt_probabilities( + halting_probs_, logits_, input, params, + batch_size_, seq_len_, feature_dim_); + + // Capture effective weights before state update + const float* p_halt = halting_probs_.host(); + const float* cum_halt = cum_halt_ptr; + const float* remainders = remainders_ptr; + float* true_weights = true_effective_weights_.host(); + + for (long pos = 0; pos < total_positions; ++pos) { + if (cum_halt[pos] < halt_threshold_) { + float p = p_halt[pos]; + float r = remainders[pos]; + + // Compute effective weight: alpha_t^n = min(p * rho, theta - h_t^(n-1)) + float effective = std::min(p * r, halt_threshold_ - cum_halt[pos]); + + // Store for backward pass + true_weights[pos] += effective; + } + } + + // Update ACT state and accumulate weighted outputs + tt::update_act_state( + output, input, halting_probs_, + cumulative_halting_, remainders_, n_steps_, + batch_size_, seq_len_, d_model_, num_channels_, + halt_threshold_, step + ); + + // Early termination optimization + if (all_positions_halted(cumulative_halting_)) break; + } + + // Finalize with remainder contributions + tt::finalize_act_output( + output, input, remainders_, + batch_size_, seq_len_, d_model_, num_channels_); + + // Add remainder weights for gradient computation + const float* final_remainders = remainders_.host(); + float* true_weights = true_effective_weights_.host(); + for (long pos = 0; pos < total_positions; ++pos) { + if (final_remainders[pos] > 1e-6f) { + true_weights[pos] += final_remainders[pos]; + } + } + + // Compute statistics for monitoring and regularization + compute_ponder_stats(); + } + + template + void backward(const tensor& gradient_input, SUBNET& sub, tensor& params_grad) { + tensor& input_grad = sub.get_gradient_input(); + + const float* grad_in = gradient_input.host(); + const float* eff_weights = true_effective_weights_.host(); + float* grad_out = input_grad.host(); + + for (long n = 0; n < batch_size_; ++n) { + for (long s = 0; s < seq_len_; ++s) { + const long pos = n * seq_len_ + s; + const float weight = eff_weights[pos]; + + for (long c = 0; c < num_channels_; ++c) { + for (long d = 0; d < d_model_; ++d) { + const long idx = ((n * num_channels_ + c) * seq_len_ + s) * d_model_ + d; + grad_out[idx] += weight * grad_in[idx]; + } + } + } + } + + // Compute parameter gradients from ponder cost regularization + params_grad = 0; + + // Optional: Apply depth-dependent gradient scaling + if (enable_depth_scaling_) { + tt::apply_act_depth_scaling( + input_grad, n_steps_, + batch_size_, seq_len_, d_model_, num_channels_, + static_cast(max_steps_), 0.1f + ); + } + } + + // Accessor methods + const tensor& get_layer_params() const { return params; } + tensor& get_layer_params() { return params; } + long get_max_steps() const { return max_steps_; } + float get_halt_threshold() const { return halt_threshold_; } + float get_ponder_penalty() const { return ponder_penalty_; } + + void set_halt_threshold(float threshold) { + if (threshold > 0 && threshold <= 1.0f) + halt_threshold_ = threshold; + } + void set_ponder_penalty(float penalty) { + if (penalty >= 0) + ponder_penalty_ = penalty; + } + + // Statistics for monitoring and regularization + float get_ponder_cost() const { return ponder_cost_; } // R(x) + float get_average_steps() const { return avg_steps_; } // Average N(t) + + // Depth scaling control + void enable_depth_scaling() { enable_depth_scaling_ = true; } + void disable_depth_scaling() { enable_depth_scaling_ = false; } + bool depth_scaling_enabled() const { return enable_depth_scaling_; } + + inline dpoint map_input_to_output(const dpoint& p) const { return p; } + inline dpoint map_output_to_input(const dpoint& p) const { return p; } + + // Serialization methods + friend void serialize(const adaptive_computation_time_& item, std::ostream& out) { + dlib::serialize("act_", out); + dlib::serialize(item.max_steps_, out); + dlib::serialize(item.halt_threshold_, out); + dlib::serialize(item.ponder_penalty_, out); + dlib::serialize(item.enable_depth_scaling_, out); + dlib::serialize(item.batch_size_, out); + dlib::serialize(item.seq_len_, out); + dlib::serialize(item.d_model_, out); + dlib::serialize(item.num_channels_, out); + dlib::serialize(item.feature_dim_, out); + dlib::serialize(item.params, out); + } + + friend void deserialize(adaptive_computation_time_& item, std::istream& in) { + std::string version; + dlib::deserialize(version, in); + if (version != "act_") + throw serialization_error("Unexpected version: " + version); + dlib::deserialize(item.max_steps_, in); + dlib::deserialize(item.halt_threshold_, in); + dlib::deserialize(item.ponder_penalty_, in); + dlib::deserialize(item.enable_depth_scaling_, in); + dlib::deserialize(item.batch_size_, in); + dlib::deserialize(item.seq_len_, in); + dlib::deserialize(item.d_model_, in); + dlib::deserialize(item.num_channels_, in); + dlib::deserialize(item.feature_dim_, in); + dlib::deserialize(item.params, in); + + item.allocate_workspace(); + } + + friend std::ostream& operator<<(std::ostream& out, const adaptive_computation_time_& item) { + out << "act (steps=" << item.max_steps_ + << ", dim=" << item.feature_dim_ + << ", threshold=" << item.halt_threshold_ + << ", penalty=" << item.ponder_penalty_ << ")"; + return out; + } + + friend void to_xml(const adaptive_computation_time_& item, std::ostream& out) { + out << "\n"; + out << mat(item.params); + out << "\n"; + } + + private: + void allocate_workspace() { + const long total_positions = batch_size_ * seq_len_; + + // Allocate state tensors for maximum expected size + // These track the ACT state for each position (batch, sequence) + halting_probs_.set_size(total_positions, 1, 1, 1); // p_t^n + cumulative_halting_.set_size(total_positions, 1, 1, 1); // h_t^n + remainders_.set_size(total_positions, 1, 1, 1); // rho_t + n_steps_.set_size(total_positions, 1, 1, 1); // N(t) + logits_.set_size(total_positions, 1, 1, 1); // logits before sigmoid + grad_logits_.set_size(total_positions, 1, 1, 1); // gradient w.r.t. logits + true_effective_weights_.set_size(total_positions, 1, 1, 1); + + // Input cache needs full dimensions for gradient computation + input_cache_.set_size(batch_size_, num_channels_, seq_len_, d_model_); + } + + bool all_positions_halted(const resizable_tensor& ch) const { + const float* cum_halt = ch.host(); + const long total = batch_size_ * seq_len_; + + for (long i = 0; i < total; ++i) { + if (cum_halt[i] < halt_threshold_) return false; + } + return true; + } + + void compute_ponder_stats() { + const float* steps = n_steps_.host(); + const long total = batch_size_ * seq_len_; + + // Compute average number of steps: (1/T) * SUM N(t) + float sum_steps = 0; + for (long i = 0; i < total; ++i) sum_steps += steps[i]; + avg_steps_ = sum_steps / total; + + // Normalize ponder cost by maximum possible steps + ponder_cost_ = avg_steps_ / max_steps_; + } + + // Configuration parameters + long max_steps_; // Maximum computation steps per position + float halt_threshold_; // theta: Halting threshold (typically 0.99) + float ponder_penalty_; // lambda: Ponder cost weight for regularization + bool enable_depth_scaling_; // Enable depth-dependent gradient scaling + + // Dimension tracking + long batch_size_; + long seq_len_; + long d_model_; + long num_channels_; + long feature_dim_; + + // Learnable parameters + resizable_tensor params; + + // Working memory + resizable_tensor halting_probs_; // p_t^n: Halting probabilities + resizable_tensor cumulative_halting_; // h_t^n: Cumulative halting probabilities + resizable_tensor remainders_; // rho_t: Remaining probability mass + resizable_tensor n_steps_; // N(t): Number of steps taken + resizable_tensor logits_; // Raw logits before sigmoid + resizable_tensor grad_logits_; // Gradients w.r.t. logits + resizable_tensor input_cache_; // Cached input for backward pass + resizable_tensor true_effective_weights_; + + // Statistics for monitoring + float ponder_cost_; // R(x): Current ponder cost + float avg_steps_; // Average number of computation steps + }; + + template + using adaptive_computation_time = add_layer, SUBNET>; + + template + using act = add_layer, SUBNET>; // Default 8 steps + + template + using act4 = add_layer, SUBNET>; // Fast version + + template + using act16 = add_layer, SUBNET>; // Deep version + // ---------------------------------------------------------------------------------------- } diff --git a/dlib/dnn/layers_abstract.h b/dlib/dnn/layers_abstract.h index 5809c7f1fe..cbfe81ad66 100644 --- a/dlib/dnn/layers_abstract.h +++ b/dlib/dnn/layers_abstract.h @@ -4699,6 +4699,125 @@ namespace dlib template using tril_diag = add_layer, SUBNET>; +// ---------------------------------------------------------------------------------------- + + template + class adaptive_computation_time_ + { + /*! + REQUIREMENTS ON TEMPLATE ARGUMENTS + - max_steps > 0 + + WHAT THIS OBJECT REPRESENTS + This is an implementation of the EXAMPLE_COMPUTATIONAL_LAYER_ interface + defined above. It implements Adaptive Computation Time (ACT) following + Graves (2016) "Adaptive Computation Time for Recurrent Neural Networks" + (arXiv:1603.08983). + + ACT allows the network to adaptively determine how many computational steps + to perform for each position in the input sequence, spending more computation + on difficult parts while quickly processing easier parts. + + MATHEMATICAL FOUNDATION: + Core ACT Algorithm: + - For each sequence position t, perform up to max_steps computational steps + - At step n, compute halting probability: p_t^n = sigmoid(W_halt^T * s_t^n + b_halt) + - Cumulative halting probability: h_t^n = sum_{i=1 to n} p_t^i + - Stop when h_t^n >= theta (threshold, typically 0.99) + - Final output: y_t = sum_{n=1 to N(t)} alpha_t^n * y_hat_t^n + + Where alpha_t^n (effective weight) is computed as: + - alpha_t^n = p_t^n * rho_t^{n-1} for intermediate steps + - alpha_t^{N(t)} = 1 - h_t^{N(t)-1} (remainder for final step) + - rho_t^n = 1 - h_t^n (remaining probability mass) + + PONDER COST (Regularization): + - R(x) = (1/T) * sum_t (N(t) + rho_t^{N(t)}) + - Total loss: L = L_task + lambda * R(x) + - lambda is controlled by get_ponder_penalty() + + IMPLEMENTATION DETAILS: + - Input/Output tensors have identical dimensions + - Learnable parameters: W_halt in R^{d x k}, b_halt in R + - State tracking per position: cumulative halting, remainders, step counts + - Early termination when all positions halt + !*/ + + public: + adaptive_computation_time_(); + /*! + ensures + - #get_max_steps() == max_steps + - #get_halt_threshold() == 0.99f + - #get_ponder_penalty() == 0.01f + !*/ + + // Configuration accessors + long get_max_steps() const; + float get_halt_threshold() const; + float get_ponder_penalty() const; + + void set_halt_threshold(float threshold); + /*! + requires + - 0 < threshold <= 1.0f + ensures + - #get_halt_threshold() == threshold + !*/ + + void set_ponder_penalty(float penalty); + /*! + requires + - penalty >= 0 + ensures + - #get_ponder_penalty() == penalty + !*/ + + // Runtime statistics + float get_ponder_cost() const; + /*! + ensures + - returns the ponder cost R(x) from the most recent forward pass + - value represents average computational cost per position + !*/ + + float get_average_steps() const; + /*! + ensures + - returns the average number of computation steps per position + - value is between 1.0 and max_steps + !*/ + + // Layer interface + template void setup(const SUBNET& sub); + template void forward(const SUBNET& sub, resizable_tensor& output); + template void backward(const tensor& gradient_input, SUBNET& sub, tensor& params_grad); + dpoint map_input_to_output(dpoint p) const; + dpoint map_output_to_input(dpoint p) const; + const tensor& get_layer_params() const; + tensor& get_layer_params(); + + friend void serialize(const adaptive_computation_time_& item, std::ostream& out); + friend void deserialize(adaptive_computation_time_& item, std::istream& in); + friend std::ostream& operator<<(std::ostream& out, const adaptive_computation_time_& item); + friend void to_xml(const adaptive_computation_time_& item, std::ostream& out); + /*! + provides serialization support and output operators + !*/ + }; + + template + using adaptive_computation_time = add_layer, SUBNET>; + + template + using act = add_layer, SUBNET>; + + template + using act4 = add_layer, SUBNET>; + + template + using act16 = add_layer, SUBNET>; + // ---------------------------------------------------------------------------------------- } diff --git a/dlib/test/dnn.cpp b/dlib/test/dnn.cpp index c564e277e1..25eafd2709 100644 --- a/dlib/test/dnn.cpp +++ b/dlib/test/dnn.cpp @@ -860,86 +860,136 @@ namespace // ---------------------------------------------------------------------------------------- -void test_positional_encodings() -{ - print_spinner(); - using net_type = tag1>>>; - net_type net; + void test_positional_encodings() + { + print_spinner(); + using net_type = tag1>>>; + net_type net; - const unsigned long sequence_dim = 4; - const unsigned long embedding_dim = 6; - const unsigned long n_samples = 1, n_channels = 1; - matrix input_data(sequence_dim, embedding_dim); - input_data = 0.0f; + const unsigned long sequence_dim = 4; + const unsigned long embedding_dim = 6; + const unsigned long n_samples = 1, n_channels = 1; + matrix input_data(sequence_dim, embedding_dim); + input_data = 0.0f; - resizable_tensor input_tensor(n_samples, n_channels, sequence_dim, embedding_dim); - std::vector> x(n_samples); - x[0] = input_data; - net.to_tensor(&x[0], &x[0] + n_samples, input_tensor); - net.forward(input_tensor); + resizable_tensor input_tensor(n_samples, n_channels, sequence_dim, embedding_dim); + std::vector> x(n_samples); + x[0] = input_data; + net.to_tensor(&x[0], &x[0] + n_samples, input_tensor); + net.forward(input_tensor); - matrix expected_output(sequence_dim, embedding_dim); - const float n = 10000.0f; - for (long r = 0; r < sequence_dim; ++r) { - for (long c = 0; c < embedding_dim; ++c) { - float theta = static_cast(r) / std::pow(n, static_cast(c) / embedding_dim); - expected_output(r, c) = (c % 2 == 0) ? std::sin(theta) : std::cos(theta); - } - } + matrix expected_output(sequence_dim, embedding_dim); + const float n = 10000.0f; + for (long r = 0; r < sequence_dim; ++r) { + for (long c = 0; c < embedding_dim; ++c) { + float theta = static_cast(r) / std::pow(n, static_cast(c) / embedding_dim); + expected_output(r, c) = (c % 2 == 0) ? std::sin(theta) : std::cos(theta); + } + } - auto& net_output = layer(net).get_output(); - DLIB_TEST(max(abs(mat(net_output) - expected_output)) < 1e-5); -} + auto& net_output = layer(net).get_output(); + DLIB_TEST(max(abs(mat(net_output) - expected_output)) < 1e-5); + } // ---------------------------------------------------------------------------------------- -void test_embeddings() -{ - print_spinner(); - const size_t num_sequences = 100, sequence_length = 7, num_classes = 3, num_tokens = 50, embedding_length = 5; - using net_type = loss_multiclass_log>>>>>>>>; - net_type net; - dnn_trainer trainer(net, sgd(0, 0.9)); - trainer.set_learning_rate(1e-1); - trainer.set_min_learning_rate(1e-4); - trainer.set_mini_batch_size(16); - trainer.set_max_num_epochs(500); - - dlib::rand rnd(std::rand()); - auto generate_sequences = [&](size_t num_sequences, size_t sequence_length, size_t num_tokens) { - std::vector> sequences; - for (size_t i = 0; i < num_sequences; ++i) - { - matrix seq(sequence_length, 1); - for (size_t j = 0; j < sequence_length; ++j) - seq(j, 0) = rnd.get_random_32bit_number() % num_tokens; - sequences.push_back(seq); - } - return sequences; - }; - - auto generate_labels = [&](size_t num_sequences, size_t num_classes) { - std::vector labels; - for (size_t i = 0; i < num_sequences; ++i) - labels.push_back(rnd.get_random_32bit_number() % num_classes); - return labels; - }; - - auto sequences = generate_sequences(num_sequences, sequence_length, num_tokens); - auto labels = generate_labels(num_sequences, num_classes); - - trainer.train(sequences, labels); - std::vector predicted_labels = net(sequences); - size_t num_correct = 0; - for (size_t i = 0; i < labels.size(); ++i) - if (predicted_labels[i] == labels[i]) ++num_correct; - - double acc = static_cast(num_correct) / labels.size(); - DLIB_TEST(acc > 0.9); -} + void test_embeddings() + { + print_spinner(); + const size_t num_sequences = 100, sequence_length = 7, num_classes = 3, num_tokens = 50, embedding_length = 5; + using net_type = loss_multiclass_log>>>>>>>>; + net_type net; + dnn_trainer trainer(net, sgd(0, 0.9)); + trainer.set_learning_rate(1e-1); + trainer.set_min_learning_rate(1e-4); + trainer.set_mini_batch_size(16); + trainer.set_max_num_epochs(500); + + dlib::rand rnd(std::rand()); + auto generate_sequences = [&](size_t num_sequences, size_t sequence_length, size_t num_tokens) { + std::vector> sequences; + for (size_t i = 0; i < num_sequences; ++i) + { + matrix seq(sequence_length, 1); + for (size_t j = 0; j < sequence_length; ++j) + seq(j, 0) = rnd.get_random_32bit_number() % num_tokens; + sequences.push_back(seq); + } + return sequences; + }; + + auto generate_labels = [&](size_t num_sequences, size_t num_classes) { + std::vector labels; + for (size_t i = 0; i < num_sequences; ++i) + labels.push_back(rnd.get_random_32bit_number() % num_classes); + return labels; + }; + + auto sequences = generate_sequences(num_sequences, sequence_length, num_tokens); + auto labels = generate_labels(num_sequences, num_classes); + + trainer.train(sequences, labels); + std::vector predicted_labels = net(sequences); + size_t num_correct = 0; + for (size_t i = 0; i < labels.size(); ++i) + if (predicted_labels[i] == labels[i]) ++num_correct; + + double acc = static_cast(num_correct) / labels.size(); + DLIB_TEST(acc > 0.9); + } + +// ---------------------------------------------------------------------------------------- + + void test_adaptive_computation_time_network() + { + print_spinner(); + + // Test with a simple network containing ACT layer + using net_type = tag1>>>; + net_type net; + + // Create test data + std::vector> training_images; + const int seq_len = 4, d_model = 3; + + matrix sample(seq_len, d_model); + dlib::rand rnd(54321); + + for (int i = 0; i < 10; ++i) + { + for (int r = 0; r < seq_len; ++r) + for (int c = 0; c < d_model; ++c) + sample(r, c) = rnd.get_random_gaussian(); + training_images.push_back(sample); + } + + // Convert to tensor and run forward pass + resizable_tensor input_tensor; + net.to_tensor(&training_images[0], &training_images[0] + training_images.size(), input_tensor); + + // Forward pass + net.forward(input_tensor); + + // Get output and verify dimensions + const tensor& output = net.get_output(); + DLIB_TEST(output.num_samples() == training_images.size()); + DLIB_TEST(output.nr() == seq_len); + DLIB_TEST(output.nc() == d_model); + + // Access the ACT layer for statistics + auto& act_layer = layer(net).subnet().layer_details(); + dlog << LINFO << "Network ACT ponder cost: " << act_layer.get_ponder_cost(); + dlog << LINFO << "Network ACT average steps: " << act_layer.get_average_steps(); + + // Verify reasonable statistics + DLIB_TEST(act_layer.get_ponder_cost() >= 0.0f && act_layer.get_ponder_cost() <= 1.0f); + DLIB_TEST(act_layer.get_average_steps() >= 1.0f); + + dlog << LINFO << "ACT network tests completed successfully"; + } // ---------------------------------------------------------------------------------------- @@ -2551,6 +2601,12 @@ void test_embeddings() auto res = test_layer(l); DLIB_TEST_MSG(res, res); } + { + print_spinner(); + adaptive_computation_time_<6> l; + auto res = test_layer(l); + DLIB_TEST_MSG(res, res); + } } // ---------------------------------------------------------------------------------------- @@ -5211,6 +5267,7 @@ void test_multm_prev() test_positional_encodings(); test_embeddings(); test_tril(); + test_adaptive_computation_time_network(); test_basic_tensor_ops(); test_resize_to(); test_layers();