Arch-LLM: Taming LLMs for Neural Architecture Generation via Unsupervised Discrete Representation Learning

📄 arXiv: 2503.22063v1 📥 PDF

作者: Deshani Geethika Poddenige, Sachith Seneviratne, Damith Senanayake, Mahesan Niranjan, PN Suganthan, Saman Halgamuge

分类: cs.LG

发布日期: 2025-03-28

💡 一句话要点

Arch-LLM：利用无监督离散表示学习，驯服LLM以生成神经架构

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱四：生成式动作 (Generative Motion) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 神经架构搜索 无监督学习 离散表示学习 变分自编码器 大型语言模型 架构生成 VQ-VAE

📋 核心要点

  1. 现有基于VAE的神经架构搜索方法，将离散架构映射到连续空间，导致生成大量无效或重复架构。
  2. 提出Arch-LLM，使用VQ-VAE学习与离散神经架构对齐的离散潜在空间，并用LLM生成架构序列。
  3. 实验表明，Arch-LLM在NAS-Bench-101和NAS-Bench-201上，显著提高了有效和唯一架构的生成比例。

📝 摘要（中文）

无监督表示学习已广泛应用于各种模态，包括神经架构，并在神经架构搜索（NAS）等下游应用中发挥着关键作用。这些方法通常学习一个无监督的表示空间，然后在下游搜索中生成/采样架构。一种常见的方法是使用变分自编码器（VAE）将离散架构映射到连续表示空间，然而，从这些空间采样通常会导致高比例的无效或重复的神经架构。这可能是由于固有的离散架构空间到连续空间的非自然映射，这强调了对这些架构的鲁棒离散表示的需求。为了解决这个问题，我们引入了一种向量量化变分自编码器（VQ-VAE），以学习与离散神经架构更自然对齐的离散潜在空间。与VAE相比，VQ-VAE（i）将每个架构映射到离散代码序列，并且（ii）允许先验由任何生成模型学习，而不是假设正态分布。然后，我们将这些架构潜在代码表示为数字序列，并利用大型语言模型训练文本到文本模型，以学习和生成表示架构的序列。我们在Inception/ResNet类基于单元的搜索空间（即NAS-Bench-101和NAS-Bench-201）上进行了实验。与基于VAE的方法相比，我们的方法在NASBench-101上将有效和唯一架构的生成提高了80%以上，在NASBench-201上提高了8%以上。最后，我们通过采用基于序列建模的NAS算法，证明了我们的方法在NAS中的适用性。

🔬 方法详解

问题定义：现有基于VAE的神经架构搜索方法，试图将离散的神经架构空间映射到连续的潜在空间。然而，这种映射方式并不自然，导致从潜在空间采样时，会产生大量无效或重复的架构，严重影响了搜索效率和质量。因此，如何学习一种能够有效表示离散神经架构的潜在空间，并从中生成高质量架构，是本文要解决的核心问题。

核心思路：本文的核心思路是利用VQ-VAE学习神经架构的离散表示，并使用大型语言模型（LLM）来生成架构序列。VQ-VAE能够将每个架构映射到离散的代码序列，从而更好地捕捉架构的离散特性。然后，将这些代码序列作为LLM的输入，训练LLM学习架构的分布，并生成新的架构序列。这种方法避免了将离散空间强行映射到连续空间的问题，从而提高了生成有效架构的概率。

技术框架：Arch-LLM的整体框架主要包含两个阶段：1）使用VQ-VAE学习神经架构的离散表示；2）使用LLM生成架构序列。在第一阶段，VQ-VAE将输入的神经架构编码为离散的代码序列，并学习一个码本，用于将编码后的向量量化为离散的码字。在第二阶段，将VQ-VAE生成的代码序列作为LLM的输入，训练LLM学习架构的分布，并使用LLM生成新的代码序列，然后通过VQ-VAE的解码器将其转换为神经架构。

关键创新：本文最重要的技术创新点在于使用VQ-VAE学习神经架构的离散表示，并结合LLM进行架构生成。与传统的VAE方法相比，VQ-VAE能够更好地捕捉架构的离散特性，从而生成更高质量的架构。此外，使用LLM作为生成模型，可以利用LLM强大的语言建模能力，学习更复杂的架构分布，从而生成更多样化的架构。

关键设计：在VQ-VAE的设计中，使用了标准的VQ-VAE结构，包括编码器、解码器和码本。码本的大小是一个重要的参数，需要根据具体的搜索空间进行调整。在LLM的设计中，使用了Transformer架构，并采用了文本到文本的训练方式。损失函数包括VQ-VAE的重构损失和码本损失，以及LLM的语言建模损失。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Arch-LLM在NAS-Bench-101上将有效和唯一架构的生成提高了80%以上，在NAS-Bench-201上提高了8%以上。与基于VAE的方法相比，Arch-LLM能够生成更高比例的有效和唯一架构，证明了其在神经架构搜索中的有效性。此外，实验还验证了Arch-LLM在NAS中的适用性，表明其可以与其他NAS算法相结合，进一步提升搜索性能。

🎯 应用场景

Arch-LLM可应用于各种神经架构搜索任务，尤其是在计算资源有限的情况下，能够高效生成高质量的候选架构。该方法还可扩展到其他离散结构的设计，例如硬件设计、药物发现等领域，具有广阔的应用前景。未来，可以探索将Arch-LLM与其他NAS算法相结合，进一步提升搜索效率和性能。

📄 摘要（原文）

Unsupervised representation learning has been widely explored across various modalities, including neural architectures, where it plays a key role in downstream applications like Neural Architecture Search (NAS). These methods typically learn an unsupervised representation space before generating/ sampling architectures for the downstream search. A common approach involves the use of Variational Autoencoders (VAEs) to map discrete architectures onto a continuous representation space, however, sampling from these spaces often leads to a high percentage of invalid or duplicate neural architectures. This could be due to the unnatural mapping of inherently discrete architectural space onto a continuous space, which emphasizes the need for a robust discrete representation of these architectures. To address this, we introduce a Vector Quantized Variational Autoencoder (VQ-VAE) to learn a discrete latent space more naturally aligned with the discrete neural architectures. In contrast to VAEs, VQ-VAEs (i) map each architecture into a discrete code sequence and (ii) allow the prior to be learned by any generative model rather than assuming a normal distribution. We then represent these architecture latent codes as numerical sequences and train a text-to-text model leveraging a Large Language Model to learn and generate sequences representing architectures. We experiment our method with Inception/ ResNet-like cell-based search spaces, namely NAS-Bench-101 and NAS-Bench-201. Compared to VAE-based methods, our approach improves the generation of valid and unique architectures by over 80% on NASBench-101 and over 8% on NASBench-201. Finally, we demonstrate the applicability of our method in NAS employing a sequence-modeling-based NAS algorithm.