Growing Transformers Modular Expansion

Growing Transformers: Progressive Model Construction on Frozen Foundations

Traditional large language model training requires end-to-end gradient flow through billions of parameters, demanding massive compute clusters and suffering from optimization instabilities. Growing Transformers inverts this approach through constructive learning: build models incrementally by stacking Transformer layers on a fixed embedding foundation, training each new layer to convergence before freezing and adding the next. This strategy reduces memory requirements, enables single-GPU training of large models, and paradoxically improves performance by ~5% compared to training entire models from scratch.

The core insight is that semantic understanding emerges from architectural depth and composition rather than sophisticated input embeddings. By freezing random or minimally-trained embeddings, you force each subsequent layer to learn meaningful transformations from the layer below, creating a strong inductive bias toward hierarchical feature extraction.

Core Concept

Growing Transformers uses constructive training—a two-component strategy:

1. Layer-wise Sequential Addition: Start with a frozen embedding layer and trained lower layers. Add one Transformer block at a time, train it to convergence in isolation, then freeze before adding the next.
2. Holistic Fine-tuning via LoRA: After building the full stack, apply Low-Rank Adaptation (LoRA) across all layers for final performance gains without retraining from scratch.

Each new layer sees identical training conditions: fixed representations from below and a standard language modeling objective. The model grows incrementally while remaining trainable on single GPUs, then benefits from whole-model fine-tuning after construction completes.

Architecture Overview

• Frozen Embedding Layer: Random or pretrained embeddings (trainable dimension: 0 or 1B parameters) that serve as the non-learnable base
• Sequential Transformer Blocks: Standard 2-layer → 3-layer → 4-layer progression, each trained independently until convergence
• LoRA Fine-tuning Head: Low-rank weight matrices applied to attention and feed-forward layers across the full stack
• Training Scheduler: Per-layer convergence detection (loss plateau) with early stopping to determine when to add next layer
• Unified Tokenizer: Shared token vocabulary throughout training (no vocabulary shifts during growth)
• Checkpoint Management: Persistent storage of frozen layers during sequential training to enable resumption

Implementation

The following demonstrates layer-by-layer training with checkpoint management:

import torch
import torch.nn as nn
from typing import Optional, List

class FrozenEmbedding(nn.Module):
    """Non-trainable embedding layer as training foundation."""
    def __init__(self, vocab_size: int, embed_dim: int, initialize_random: bool = True):
        super().__init__()
        if initialize_random:
            self.embed = nn.Embedding(vocab_size, embed_dim)
            # Freeze immediately—no training signal flows through
            for param in self.embed.parameters():
                param.requires_grad = False
        else:
            self.embed = nn.Embedding(vocab_size, embed_dim)

    def forward(self, input_ids):
        return self.embed(input_ids)

class TransformerBlock(nn.Module):
    """Single Transformer layer for independent training."""
    def __init__(self, hidden_dim: int, num_heads: int, ff_dim: int, dropout: float = 0.1):
        super().__init__()
        self.attention = nn.MultiheadAttention(hidden_dim, num_heads, dropout=dropout, batch_first=True)
        self.norm1 = nn.LayerNorm(hidden_dim)
        self.norm2 = nn.LayerNorm(hidden_dim)

        self.feed_forward = nn.Sequential(
            nn.Linear(hidden_dim, ff_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ff_dim, hidden_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, attn_mask=None):
        # Self-attention with residual
        attn_out, _ = self.attention(x, x, x, attn_mask=attn_mask)
        x = x + attn_out
        x = self.norm1(x)

        # Feed-forward with residual
        ff_out = self.feed_forward(x)
        x = x + ff_out
        x = self.norm2(x)
        return x

class GrowingTransformer(nn.Module):
    """Model that grows layer-by-layer on frozen embeddings."""
    def __init__(self, vocab_size: int, hidden_dim: int, num_heads: int, ff_dim: int):
        super().__init__()
        self.embedding = FrozenEmbedding(vocab_size, hidden_dim, initialize_random=True)
        self.blocks = nn.ModuleList()  # Grows during training
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.ff_dim = ff_dim

        # Initialize with first layer
        self.add_block()

    def add_block(self):
        """Add a new trainable Transformer layer."""
        new_block = TransformerBlock(self.hidden_dim, self.num_heads, self.ff_dim)
        self.blocks.append(new_block)

    def forward(self, input_ids):
        x = self.embedding(input_ids)
        for block in self.blocks:
            x = block(x)
        return x

def train_single_layer(model: GrowingTransformer, layer_idx: int, train_loader,
                       optimizer, criterion, num_epochs: int = 100,
                       patience: int = 5) -> float:
    """Train a single layer while freezing all others."""
    # Freeze all blocks except the one being trained
    for i, block in enumerate(model.blocks):
        for param in block.parameters():
            param.requires_grad = (i == layer_idx)

    best_loss = float('inf')
    patience_counter = 0

    for epoch in range(num_epochs):
        total_loss = 0
        for batch_idx, (input_ids, target_ids) in enumerate(train_loader):
            optimizer.zero_grad()

            # Forward pass through all layers (but only latest layer trains)
            logits = model(input_ids)
            logits = logits.view(-1, model.hidden_dim)  # Flatten for language modeling head
            target = target_ids.view(-1)

            loss = criterion(logits, target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        avg_loss = total_loss / len(train_loader)

        if avg_loss < best_loss:
            best_loss = avg_loss
            patience_counter = 0
        else:
            patience_counter += 1

        if patience_counter >= patience:
            print(f"Layer {layer_idx} converged at epoch {epoch} with loss {best_loss:.4f}")
            break

    return best_loss

def sequential_growth(initial_model: GrowingTransformer, train_loader,
                      num_growth_steps: int = 5, epochs_per_layer: int = 100):
    """Grow model by repeatedly adding and training layers."""
    for step in range(num_growth_steps):
        print(f"\n=== Growth Step {step + 1}/{num_growth_steps} ===")

        # Add new layer
        if step > 0:  # First layer already exists
            initial_model.add_block()

        # Train newly added layer
        optimizer = torch.optim.AdamW(initial_model.parameters(), lr=1e-4)
        criterion = nn.CrossEntropyLoss()

        layer_loss = train_single_layer(
            initial_model,
            layer_idx=step,
            train_loader=train_loader,
            optimizer=optimizer,
            criterion=criterion,
            num_epochs=epochs_per_layer,
            patience=5
        )

        # Freeze current layer before proceeding
        for param in initial_model.blocks[step].parameters():
            param.requires_grad = False

        print(f"Layer {step} frozen. Model now has {len(initial_model.blocks)} layers")

    return initial_model

# Example usage
def apply_lora_finetuning(model: GrowingTransformer, train_loader, lr: float = 1e-3, epochs: int = 10):
    """Apply LoRA fine-tuning across all layers after growth."""
    # Add LoRA parameters to attention layers
    for block in model.blocks:
        # Simplified: in practice use peft library for proper LoRA
        block.lora_q = nn.Linear(model.hidden_dim, 8)  # Low-rank projection
        block.lora_v = nn.Linear(8, model.hidden_dim)

    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=lr
    )

    for epoch in range(epochs):
        total_loss = 0
        for input_ids, target_ids in train_loader:
            optimizer.zero_grad()
            logits = model(input_ids)
            loss = nn.CrossEntropyLoss()(logits.view(-1, model.hidden_dim), target_ids.view(-1))
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        print(f"LoRA Fine-tuning Epoch {epoch + 1}: Loss {total_loss / len(train_loader):.4f}")

This implementation demonstrates the core constructive learning pipeline: sequentially adding layers, training each in isolation, freezing, and then fine-tuning the complete stack.

Practical Guidance

When to Use Growing Transformers

• Resource-constrained training: Fitting billion-parameter models on single GPUs
• Incremental model deployment: Adding capacity to existing models without retraining from scratch
• Efficient scaling experiments: Quickly exploring 1B, 2B, 3B variants with minimal overhead
• Distributed training pipelines: Different layers can train on different GPUs without synchronization
• Research on layer dynamics: Analyzing how representational capacity evolves during sequential growth

When NOT to Use

• State-of-the-art performance: End-to-end training still achieves slightly higher performance on massive datasets
• Models <100M parameters: Overhead of layer management exceeds benefits; use standard training
• Recurrent or dynamic architectures: Constructive learning assumes layer independence; doesn't apply to RNNs or adaptive models
• Multi-task or continual learning with shared task-specific heads: Frozen embeddings may bottleneck task adaptation
• Tight inference latency requirements: Layer management and checkpoint switching add minimal but measurable overhead

Common Pitfalls

1. Embedding Dimension Too Small: Frozen embeddings with insufficient capacity become a bottleneck. Use hidden_dim >= 512 for models >1B parameters.
2. Training New Layers Too Long: Overtraining a single layer before adding the next reduces diversity. Use early stopping aggressively.
3. Skipping LoRA Fine-tuning: Layer-wise training creates suboptimal inter-layer interactions. Always perform end-to-end LoRA fine-tuning afterward.
4. Learning Rate Schedule Mismatch: Using identical LR for all layers causes instability. Freeze layers with larger learning rates to stabilize training.
5. Ignoring Validation Performance: Monitor validation loss during layer training to detect when adding the next layer is premature.

Reference

Ke, Y., Saha, S., et al. (2025). Growing Transformers: Modular Composition and Layer-wise Expansion on a Frozen Substrate. arXiv preprint arXiv:2507.07129.

Available at: https://arxiv.org/abs/2507.07129