💻 computer science

TAS-LoRA: Transformer Architecture Search with Mixture-of-LoRA Experts

TAS-LoRA is a novel Transformer architecture search method that addresses the feature collapse problem in existing approaches by introducing a Mixture-of-LoRA Experts strategy with dynamic routing and group-wise initialization to enable subnet-specific feature learning while maintaining computational efficiency.

Original authors: Jeimin Jeon, Hyunju Lee, Bumsub Ham

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jeimin Jeon, Hyunju Lee, Bumsub Ham

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build the perfect car for every possible driver on the road: a tiny city commuter, a heavy-duty trucker, and a luxury racer.

In the world of Artificial Intelligence (AI), specifically Vision Transformers (ViTs), these "cars" are neural networks designed to recognize images. Usually, engineers manually design these networks, which is slow and expensive. To speed this up, researchers use a method called Transformer Architecture Search (TAS).

Think of TAS as building a massive "Supernet"—a giant, all-in-one factory that contains every possible car design inside it. Instead of building each car separately, the factory shares the same engine parts (weights) for all the cars. The idea is to train this one giant factory once, and then just "pull out" the specific car design you need.

The Problem: The "One-Size-Fits-All" Trap

The paper identifies a major flaw in this factory approach, which they call "Feature Collapse."

Because every car design in the factory shares the exact same engine parts, they all end up driving the same way. The factory learns a "generic" driving style that works okay for everyone, but it's not perfect for anyone.

The city commuter needs to be agile.
The truck needs to be strong.
The racer needs to be fast.

But because they all share the same parts, the factory forces them all to drive like a "middle-of-the-road" sedan. None of them reach their full potential. In AI terms, the different network designs fail to learn the specific features they need to be good at their specific jobs.

The Solution: TAS-LoRA (The "Custom Tuning Kit")

The authors propose a new method called TAS-LoRA. Instead of forcing every car to use the exact same engine settings, they add a "Custom Tuning Kit" to each car.

They use a technique called LoRA (Low-Rank Adaptation). Imagine LoRA as a small, lightweight add-on module that you can clip onto the shared engine.

The main engine (the shared weights) stays the same for everyone.
But each car gets its own unique LoRA kit that tweaks the engine just for its specific needs.

This allows the city car to learn agility, the truck to learn strength, and the racer to learn speed, all while still using the same main factory engine.

The Smart Manager: Mixture-of-LoRA-Experts (MoLE)

Here is the tricky part: There are millions of possible car designs. If you tried to give every single design its own unique tuning kit, the factory would become too big and slow to manage.

To solve this, TAS-LoRA introduces a Mixture-of-LoRA-Experts (MoLE).

Instead of millions of unique kits, the factory has a limited set of "Expert Kits" (e.g., a "Speed Kit," a "Strength Kit," an "Agility Kit").
A Smart Manager (Router) stands at the assembly line. As each car design comes down the line, the Manager looks at its specs (how many wheels, how big the engine is) and instantly decides: "This one needs the Speed Kit mixed with the Agility Kit."

This way, the factory stays small and efficient, but every car still gets a custom-tuned experience.

The "Group" Trick: Avoiding Confusion

The paper noticed a problem: When the factory first starts, the Smart Manager doesn't know which kit goes to which car. It just gives everyone the same kit, leading back to the "Feature Collapse" problem.

To fix this, they invented Group-wise Router Initialization.

Before the factory opens, they group the cars by their basic shape (e.g., all "small cars" go in Group A, all "big trucks" in Group B).
They force the Manager to give the "Speed Kit" to Group A and the "Strength Kit" to Group B right from day one.
This ensures that from the very first day of training, the different groups learn different things, preventing them from all becoming the same "generic" car.

The Results

The authors tested this on a massive dataset called ImageNet (a huge library of photos) and several other smaller datasets (like flowers, cars, and birds).

Better Performance: Their method (TAS-LoRA) produced AI models that were significantly more accurate than previous methods.
No Extra Cost: Because the "tuning kits" can be merged back into the main engine before the car is used, there is no extra cost when the AI actually looks at a picture. It runs just as fast as the old methods.
Flexibility: Unlike other methods that might need to be rebuilt if you change the hardware rules (like making the car smaller), TAS-LoRA can adapt to new rules without starting over.

Summary

In short, the paper says: "Don't force every AI design to share the exact same brain. Give them a shared brain, but let them wear different 'smart glasses' (LoRA experts) so they can see the world in the way that's best for their specific job." This makes the AI search process faster, smarter, and more effective.

Technical Summary: TAS-LoRA

Problem Statement

Transformer Architecture Search (TAS) aims to automatically discover optimal Vision Transformer (ViT) architectures, reducing the manual effort required for design. Current state-of-the-art TAS methods primarily rely on one-shot approaches, which train a "supernet" containing all possible subnets. In these weight-entangled supernets, all subnets share a common set of weights.

The paper identifies a critical limitation in this paradigm: the feature collapse problem. Because the supernet is optimized to minimize a global loss across all subnets, the shared weights tend to learn generic features that work adequately for diverse architectures but fail to capture subnet-specific characteristics. Consequently, individual subnets cannot fully leverage their unique architectural properties (e.g., specific numbers of attention heads or MLP ratios), leading to suboptimal performance compared to subnets trained independently from scratch. Empirical analysis in the paper confirms that subnets within weight-entangled supernets exhibit near-identical feature representations (high cosine similarity), whereas independently trained subnets develop distinct representations.

Methodology: TAS-LoRA

To address feature collapse while maintaining computational efficiency, the authors propose TAS-LoRA (Transformer Architecture Search with Mixture-of-LoRA Experts). The method integrates Low-Rank Adaptation (LoRA) into the supernet training process to enable specialized feature learning for each subnet without incurring the prohibitive cost of assigning unique parameters to every possible subnet configuration.

1. Mixture-of-LoRA-Experts (MoLE)

Instead of assigning a unique LoRA module to every subnet (which is computationally infeasible due to the vast search space), TAS-LoRA employs a Mixture-of-LoRA-Experts (MoLE) strategy.

Architecture: LoRA experts are applied to the linear layers of the supernet (including QKV projection and MLP layers).
Dynamic Routing: A lightweight router dynamically selects and combines a limited set of LoRA experts based on the specific architecture of the sampled subnet.
Inference Efficiency: At inference time, the selected LoRA parameters are merged into the original supernet weights. Since the subnet configuration is fixed during inference, the router runs only once, and the merged weights introduce no additional computational overhead.

2. Router Design

The router is designed to capture the sequential dependencies of transformer blocks:

Input: It takes both block-level attributes (e.g., number of attention heads, MLP ratio) and subnet-level attributes (e.g., embedding dimension, depth) as input.
Mechanism: These attributes are encoded by a learnable embedding layer and processed by an LSTM to capture sequential dependencies across blocks. The LSTM outputs hidden representations, which are passed through a fully connected "router head" to generate a probability distribution over the available LoRA experts for each layer.

3. Group-wise Router Initialization

A key challenge identified is that during early training, the router tends to assign similar weights to all experts, leading to redundancy and preventing specialization. To mitigate this, the authors introduce a group-wise router initialization technique:

Grouping: Transformer blocks across different subnets are grouped based on architectural similarities (specifically attention head count, MLP ratio, and embedding dimension).
Initialization: The number of LoRA experts is set to match the number of groups. The router is initialized such that each expert is biased toward a distinct group. This is achieved by setting specific bias terms in the router head to a positive value $\beta$ for the corresponding group-expert pair, while zeroing out others.
Effect: This forces the router to assign distinct experts to different architectural groups early in training, encouraging diverse feature learning. As training progresses, the router dynamically updates these assignments based on the subnet's properties.

Key Contributions

Identification of Feature Collapse in TAS: The paper formally identifies and analyzes the feature collapse problem in weight-entangled supernets, demonstrating that shared weights prevent subnets from learning specialized representations.
TAS-LoRA Framework: A novel method that combines LoRA with a dynamic routing mechanism (MoLE) to learn subnet-specific features efficiently. It allows subnets to refine shared features based on their architectural properties without retraining the entire supernet from scratch.
Group-wise Router Initialization: A technique to prevent expert redundancy during the early stages of training, ensuring that LoRA experts learn distinct feature representations.
State-of-the-Art Performance: Extensive experiments demonstrate that TAS-LoRA significantly outperforms existing TAS methods (e.g., AutoFormer, DYNAS, FocusFormer) across various search spaces and datasets.

Experimental Results

The authors evaluated TAS-LoRA on the ImageNet dataset and several transfer learning benchmarks (CIFAR-10/100, Flowers, CARS, INAT-19).

ImageNet Performance: TAS-LoRA consistently improved the top-1 accuracy of supernets trained with AutoFormer and DYNAS across Tiny, Small, and Base search spaces.
- For example, in the AutoFormer-T space, TAS-LoRA achieved 75.7% accuracy, outperforming the original AutoFormer (74.7%) and other SOTA methods like PreNAS-Ti (77.1% but with less flexibility).
- In the AutoFormer-B space, TAS-LoRA-B achieved 82.6% accuracy with 54M parameters, surpassing AutoFormer-B (82.4%) and matching or exceeding larger models like GLiT-B (82.3% with 96M parameters).
Transfer Learning: TAS-LoRA demonstrated superior generalization, achieving competitive or state-of-the-art results on CIFAR-10/100, Flowers, CARS, and INAT-19.
Ablation Studies:
- Initialization: Group-wise initialization consistently outperformed random initialization, confirming its role in reducing expert redundancy.
- Feature Similarity: Cosine similarity analysis showed that TAS-LoRA with group-wise initialization produced significantly more diverse feature representations among experts compared to random initialization or standard weight-entangled methods.
- Comparison with Fine-tuning: TAS-LoRA outperformed a fully fine-tuned supernet (where all weights are updated) despite training only for 50 epochs on the MoLE layers. This indicates the gains stem from the subnet-specific feature learning capability of MoLE, not merely extended training.

Significance and Claims

The paper claims that TAS-LoRA represents a significant advancement in Transformer Architecture Search by effectively solving the feature collapse problem inherent in weight-entangled supernets.

Efficiency vs. Performance: Unlike methods that require training multiple supernets (e.g., MoS) or retraining from scratch when hardware constraints change (e.g., PreNAS), TAS-LoRA trains a single supernet with subnet-specific adaptation layers. This allows for flexible adaptation to new hardware constraints (e.g., FLOPs, latency) without retraining the entire model.
Scalability: By using a limited set of experts and a dynamic router, the method scales to supernets with enormous numbers of possible subnet configurations without a linear increase in parameters or computational cost.
Generalizability: The approach is presented as a general solution applicable to any weight-entangled supernet, improving performance regardless of the underlying training strategy (e.g., AutoFormer or DYNAS).

The authors conclude that TAS-LoRA successfully mitigates feature collapse, enabling subnets to learn specialized representations and achieving state-of-the-art performance in both architecture search and transfer learning tasks.