Low-Rank GHN Decoder

• The paper demonstrates that low-rank factorization in GHN decoders reduces parameter scaling from O(d^3) to O(d^2), enabling efficient prediction of transformer-scale models.
• The decoder predicts weight matrix factors U and V without tiling, avoiding repetitive patterns and maintaining high initialization diversity.
• Integration within LoGAH yields a memory-efficient setup with improved downstream performance, achieving substantial accuracy gains with minimal hypernetwork parameters.

A Low-Rank GHN Decoder is a neural architecture component enabling Graph HyperNetworks (GHNs) to predict the parameters of extremely wide and deep target models without incurring prohibitive parameter cost or repetitive tiling artifacts. Introduced in the context of LoGAH—Low-rank GrAph Hypernetworks—a low-rank decoder replaces the conventional full-rank block generation step in GHNs with a compact, low-rank factorization, supporting the prediction of transformer-scale models (hundreds of millions of parameters) with orders of magnitude fewer parameters in the hypernetwork itself (Zhou et al., 25 May 2024). This technique is central to making direct, memory-efficient initialization of very large neural networks feasible via one-shot parameter prediction.

1. Motivation: GHN Scalability and the Need for Low-Rank Decoding

Early GHNs (e.g., GHN-3) used multi-layer perceptron (MLP) decoders to predict each weight block in a target network directly from node embeddings in the computational graph. For models with large widths or channel counts, this approach necessitated two problematic design choices:

• The decoder parameter count scales as $\mathcal{O}(d^3)$ in the GHN hidden dimension $d$, quickly becoming intractable for transformer-scale models.
• Weight tensors for wide target layers were constructed by copying and tiling small MLP-decoded blocks, resulting in repeated patterns that limit initialization diversity and can hinder downstream fine-tuning.

LoGAH addresses these limitations by employing a low-rank factorization in the decoder, eliminating block tiling and reducing parameter growth to $\mathcal{O}(d^2)$. This enables prediction of much larger networks while enhancing the diversity of the generated weights (Zhou et al., 25 May 2024).

2. Mathematical Structure of the Low-Rank Decoder

Formally, let $$W \in \mathbb{R}^{d_\mathrm{out} \times d_\mathrm{in}}$$ denote a weight matrix to be predicted. Instead of directly generating $W$, the decoder predicts a pair of low-rank factors $U \in \mathbb{R}^{d_\mathrm{out} \times r}$ and $V \in \mathbb{R}^{d_\mathrm{in} \times r}$ such that

$W = U V^T,$

where $r \ll \min(d_\mathrm{out}, d_\mathrm{in})$ is the rank hyperparameter. This reduces the output dimensionality from $d_\mathrm{out} \cdot d_\mathrm{in}$ to $r(d_\mathrm{out} + d_\mathrm{in})$.

For convolutional weight tensors $$W \in \mathbb{R}^{C_\mathrm{out} \times C_\mathrm{in} \times h \times w}$$, the tensor is first reshaped into a matrix of shape $(C_\mathrm{out} h) \times (C_\mathrm{in} w)$ before applying the same decomposition. In LoGAH, $r$ typically scales linearly with $d$ (e.g., $r \approx d/2$), balancing decoder size and representational expressivity.

3. Algorithmic Pipeline: From Graph Representation to Weight Tensors

Given a computational graph $f^G = (V, E)$ for the target neural network:

• Node embeddings $h_i^{(L)} \in \mathbb{R}^d$ are produced for each node $i$ after $L$ layers of a Graphormer.
• For each node $i$:
  1. The embedding $h_i^{(L)}$ is passed through a four-layer MLP with ReLU, yielding a matrix $\widetilde{W}_i \in \mathbb{R}^{2K \times r}$, where $K = \max(C_\mathrm{out} h, C_\mathrm{in} w)$.
  2. $\widetilde{W}_i$ is split into two $K \times r$ matrices, $A_i$ and $B_i^T$.
  3. The first $C_\mathrm{out} h$ rows of $A_i$ and the first $C_\mathrm{in} w$ columns of $B_i^T$ yield $A_i^{\mathrm{slice}}$ and $(B_i^{\mathrm{slice}})^T$.
  4. The product $A_i^{\mathrm{slice}} (B_i^{\mathrm{slice}})^T$ reconstructs the flattened weight block $W_i^{\mathrm{flat}}$.
  5. $W_i^{\mathrm{flat}}$ is reshaped back to the original $$C_\mathrm{out} \times C_\mathrm{in} \times h \times w$$ dimensions.

This approach enables the decoder to generate the full target tensor directly, obviating tiling and supporting arbitrarily large layers.

4. Parameter Complexity: Low-Rank versus Full MLP Decoder

The parameter count for the traditional GHN-3 full decoder is:

$$\#\mathrm{Param}_{\mathrm{GHN3\_dec}} = 4d^2 (16 \times 16) + 32d^2 + 8d^3 + d \cdot \mathrm{num\_class} \in \mathcal{O}(d^3).$$

For the LoGAH low-rank decoder with $r \approx d/2$:

$$\#\mathrm{Param}_{\mathrm{LoGAH\_dec}} = 4d^2 + 32d^2 + 8d (2r^2) + r K \in \mathcal{O}(d^2).$$

With practical values ($d = 64, 128, 256$), LoGAH’s parameter count is one to two orders of magnitude smaller, requiring only about $1\%$ of the full decoder's parameters for predicting 774-million-parameter models such as GPT-2 Large, while running on a single modern GPU (Zhou et al., 25 May 2024).

5. Integration within LoGAH: Predicting Large-Scale Transformers

The LoGAH protocol incorporating the low-rank decoder consists of:

1. Sampling small transformers (e.g., ViTs-1K, GPTs-1K).

2. Converting the model to a graph representation and embedding node operations.
3. Processing node embeddings through multi-layer Graphormers.
4. Using the low-rank decoder to generate weight tensors for each node.
5. Assembling these into the full parameter set $w_\mathrm{pred}$ for the target model.

This integration yields memory efficiency and supports model widths well beyond $d$ without increasing the GHN parameter count. Notably, LoGAH-Tiny with only 2.5M parameters can predict all 774M parameters of GPT-2 Large in a single pass (Zhou et al., 25 May 2024).

6. Comparative Analysis with Previous GHN Decoders

In comparison with prior GHN-3 decoders, the low-rank design in LoGAH offers:

• Reduced memory footprint: $\mathcal{O}(d^2)$ scaling, yielding a $5\times$–$10\times$ reduction in hypernetwork size for equivalent hidden widths.
• Enhanced expressivity and diversity: By removing copy-and-tile, generated weight patterns exhibit higher inter-tensor cosine distances and greater flexibility, empirically resulting in improved downstream task performance for both vision (e.g., ViT on CIFAR/ImageNet) and language tasks (e.g., GPT-2 on WikiText).
• Superior initialization performance: LoGAH achieves faster convergence and higher final accuracy compared to GHN-3, with observed 6–7% top-1 accuracy improvements on tasks such as ViT-Small initialization on CIFAR-100, despite using just $1/100$ of the parameters in the decoder (Zhou et al., 25 May 2024).

In summary, the Low-Rank GHN Decoder is a central advancement in the architectural design of LoGAH, enabling parameter-efficient, non-repetitive, and scalable one-shot prediction for large neural networks.