Parameter Efficiency and Sharing Techniques

The standard Transformer architecture, especially when scaled to large sizes, demands substantial computational resources and memory. The quadratic complexity of self-attention is one bottleneck, addressed by techniques like sparse or linear attention. However, another significant factor contributing to resource requirements is the sheer number of parameters, particularly in deep models with many layers. Techniques designed to reduce the parameter count make models smaller, potentially faster to train (per epoch, due to less data movement), and easier to deploy, often with minimal impact on performance.

The central idea revolves around parameter sharing, where the same set of weights is used in multiple parts of the network. This contrasts with standard deep networks where each layer typically has its own unique set of parameters.

Factorized Embedding Parameterization

In many Natural Language Processing (NLP) tasks, the vocabulary size ($V$) can be very large (tens or hundreds of thousands). The input embedding layer maps these discrete tokens to dense vectors of hidden size $H$. This results in an embedding matrix of size $V \times H$. In large models like BERT, the hidden size $H$ is also large (e.g., 768 or 1024). Consequently, the embedding matrix constitutes a significant portion of the total parameters. For instance, with $V=50,000$ and $H=1024$, the embedding matrix alone has over 50 million parameters.

The core insight behind factorized embedding parameterization, prominently featured in models like ALBERT (A Lite BERT for Self-supervised Learning of Language Representations), is that the embeddings serve two purposes: capturing context-independent token representations and projecting these into the hidden space of the Transformer layers. ALBERT argues that the context-independent representation might not need the full dimensionality $H$.

Instead of directly learning a $V \times H$ matrix, this technique factors it into two smaller matrices:

1. A lower-dimensional embedding matrix of size $V \times E$, where $E$ is the embedding dimension and $E \ll H$.
2. A projection matrix of size $E \times H$ that projects the lower-dimensional embeddings ($E$) up to the expected hidden dimension ($H$) required by the Transformer layers.

The total number of parameters for the embedding becomes $V \times E + E \times H$. Compared to the original $V \times H$, this can lead to substantial savings if $E$ is chosen significantly smaller than $H$. For example, if $V=50,000$, $H=1024$, and we choose $E=128$, the parameters change from $50,000 \times 1024 \approx 51.2M$ to $(50,000 \times 128) + (128 \times 1024) = 6.4M + 0.13M \approx 6.53M$. This is nearly an 8x reduction in embedding parameters.

This factorization decouples the vocabulary size from the hidden size of the Transformer, allowing for large vocabularies without disproportionately increasing the model size due to embeddings. It implicitly assumes that the semantic richness of individual words can be captured in a lower-dimensional space ($E$) before being projected into the higher-dimensional space ($H$) used for context-dependent processing within the Transformer layers.

Cross-Layer Parameter Sharing

Another effective technique, also central to ALBERT, is sharing parameters across the Transformer layers. In a standard Transformer encoder stack with $L$ layers, each layer typically has its own unique weight matrices for the multi-head self-attention (Q, K, V projections, output projection) and the position-wise feed-forward network (FFN). This means the parameters scale linearly with the number of layers.

Cross-layer parameter sharing breaks this assumption. Instead of learning $L$ distinct sets of layer parameters, a single set of parameters is learned and reused across all $L$ layers. Common variants include:

• Sharing all parameters: The exact same weights for attention and FFN sub-layers are used in every layer. This is the most aggressive form and leads to the largest parameter reduction.
• Sharing only FFN parameters: The attention parameters are unique per layer, but the FFN parameters are shared.
• Sharing only attention parameters: The FFN parameters are unique per layer, but the attention parameters are shared.

Sharing all parameters, as done in ALBERT, dramatically reduces the number of parameters related to the Transformer layers. If a single layer has $P_{layer}$ parameters, a standard $L$-layer model has $L \times P_{layer}$ parameters in the layers, whereas a fully shared model has only $P_{layer}$ parameters for all layers combined (excluding embeddings and normalization layers, which might or might not be shared).

Comparison of approximate parameter counts for embeddings and encoder layers in a standard BERT-Base model versus an ALBERT-Base model implementing factorization (E=128) and full cross-layer sharing. Note the significant reduction in both components for ALBERT.

Empirical results, particularly from the ALBERT paper, showed that cross-layer parameter sharing can yield models that are significantly smaller but perform comparably to, or sometimes even better than, their non-shared counterparts on various NLP benchmarks, especially when combined with factorized embeddings. It seems to act as a form of regularization, potentially improving generalization by preventing later layers from diverging too much in function from earlier ones. However, it can sometimes slow down training convergence compared to non-shared models, possibly because the shared parameters have to accommodate the functional requirements of different processing stages within the network depth.

Other Sharing Techniques

• Weight Tying: A common practice, especially in sequence-to-sequence models or language models, is to tie the weights of the input embedding layer and the final output linear layer that projects the decoder's output back to the vocabulary space. Since both layers operate between the token space and the model's hidden representation space, sharing these weights reduces parameters and can improve performance by enforcing a symmetry. This is compatible with factorized embeddings; if factorization is used ($V \times E$ and $E \times H$), the $V \times E$ matrix can be tied with the final projection matrix (which would then effectively be an $H \times E$ matrix followed by an $E \times V$ matrix, sharing the $V \times E$ part).

Trade-offs

Parameter sharing techniques offer compelling advantages:

• Reduced Model Size: Fewer parameters mean smaller model files and lower memory requirements for storage and deployment.
• Potentially Faster Training (per epoch): While convergence might take more steps, each epoch can be faster due to reduced communication overhead in distributed training or less memory pressure.
• Regularization Effect: Sharing can stabilize training and sometimes improve generalization.

However, there are potential drawbacks:

• Representational Capacity: Aggressive sharing might limit the model's ability to learn highly specialized functions in different layers.
• Convergence Speed: Training might require more steps or different optimization hyperparameters (like learning rates) to converge effectively compared to non-shared models.

Choosing the right parameter efficiency technique involves balancing these trade-offs. Factorized embeddings are often a relatively safe way to reduce parameters with minimal performance impact, especially for large vocabularies. Cross-layer parameter sharing is more aggressive but has proven surprisingly effective, as demonstrated by ALBERT. Evaluating these techniques empirically on the specific task and dataset remains an important step in model development. These methods contribute significantly to the ongoing effort to build powerful yet practical Transformer models.