Exponentially Increasing the Capacity-to-Computation Ratio for Conditional Computation in Deep Learning

Many state-of-the-art results obtained with deep networks are achieved with the largest models that could be trained, and if more computation power was available, we might be able to exploit much larger datasets in order to improve generalization ability. Whereas in learning algorithms such as decision trees the ratio of capacity (e.g., the number of parameters) to computation is very favorable (up to exponentially more parameters than computation), the ratio is essentially 1 for deep neural networks. Conditional computation has been proposed as a way to increase the capacity of a deep neural network without increasing the amount of computation required, by activating some parameters and computation “on-demand”, on a per-example basis. In this note, we propose a novel parametrization of weight matrices in neural networks which has the potential to increase up to exponentially the ratio of the number of parameters to computation. The proposed approach is based on turning on some parameters (weight matrices) when specific bit patterns of hidden unit activations are obtained. In order to better control for the overfitting that might result, we propose a parametrization that is tree-structured, where each node of the tree corresponds to a prefix of a sequence of sign bits, or gating units, associated with hidden units.

💡 Research Summary

The paper addresses a fundamental limitation of modern deep neural networks: the capacity‑to‑computation ratio is essentially one, meaning that every parameter is touched for each training example. In contrast, models such as decision trees can activate only a small subset of parameters for a given input, achieving an exponential disparity between the number of parameters that could be stored and the amount of computation actually performed. The authors propose a novel parametrization of weight matrices that brings this exponential advantage to deep learning through conditional computation.

The core idea is to derive a binary gating vector g of length k from each input x. This gating vector can be obtained deterministically (e.g., thresholding selected input dimensions) or stochastically (e.g., sampling each bit from a Bernoulli distribution whose mean is a sigmoid of a linear projection of x). The bits of g are then used as an index to select among up to 2^k different weight sub‑matrices. For each output unit j, a subset S_j(g) of at most k bits is extracted, and a function F_j maps this k-bit pattern to a weight vector w_j ∈ ℝ^q. In the simplest implementation, F_j is a table lookup, yielding a total parameter count of O(2^k·p·q) for a layer with input dimension p and output dimension q. Crucially, the computational cost of a forward pass remains O(p·q) – the same as in a conventional dense layer – plus a modest O(k·q) overhead for the table lookup or bit‑selection logic.

To mitigate the risk of severe over‑fitting (since each of the 2^k weight vectors would be updated only on a fraction of the training data), the authors introduce a tree‑structured “prefix‑sum” regularization. They imagine a binary tree of depth k+1 where each node corresponds to a prefix of the gating bits. Each node stores a weight matrix T(j, prefix). The final weight vector for unit j is obtained by summing the matrices along the path from the root (empty prefix) to the leaf defined by the full k-bit pattern:

 F_j(b) = Σ_{l=0}^{k} T(j, b_1…b_l).

Shorter prefixes are visited far more often, thus receiving many gradient updates and acting as strong regularizers, while deeper nodes (longer prefixes) are rarely activated and only provide fine‑grained corrections. The total number of independent parameters is still O(2^k·p·q), but effective degrees of freedom grow only as O(k·p·q). Regularization (L1 or L2 weight decay) is applied only to the parameters that are active at a given step, and a time‑difference Δt is used to compensate for the decay that would have accumulated during idle periods.

The paper also discusses the credit‑assignment problem for the gating decisions. Three possible learning strategies are outlined: (1) a variance‑reduced REINFORCE estimator (e.g., using a baseline or reinforcement learning tricks), (2) a straight‑through estimator that back‑propagates gradients through a continuous relaxation of the binary thresholds, and (3) a “noisy rectifier” approach where the contribution of each weight vector is modulated by a smooth function of the gating activations (e.g., a tanh‑based mask). The authors hypothesize that even a naïve approach—ignoring gradients w.r.t. the gating bits—might be sufficient if the gating function partitions the input space reasonably well, because the weight tables themselves can still be trained to minimize the loss.

Overall, the proposed framework promises three major benefits: (i) an exponential increase in model capacity without a corresponding increase in arithmetic operations, (ii) a principled regularization scheme that leverages the hierarchical structure of the gating bits, and (iii) flexibility in how the gating network is trained. However, the paper is largely theoretical; it lacks empirical validation on large‑scale datasets, and practical concerns such as memory consumption for the 2^k tables, GPU‑friendly implementations, and the stability of stochastic gating remain open questions.

In conclusion, the authors argue that conditional computation, when combined with a tree‑structured prefix‑sum parametrization, could be a key ingredient for scaling deep learning beyond the current hardware‑limited paradigm. Future work should focus on large‑scale experiments (e.g., speech or language corpora with billions of examples), efficient hardware implementations, automated tuning of the gating depth k, and more robust training algorithms for the gating network. If successful, this line of research would enable models that are far richer than today’s networks while keeping inference and training costs comparable to existing architectures.