Re Bottleneck Latent Restructuring

Re-Bottleneck: Restructure Frozen Autoencoder Latents

Pretrained audio autoencoders learn latent representations that compress audio into semantic vectors, but these latents often lack interpretable structure. Audio researchers frequently need to impose specific constraints—ensuring channels capture frequencies in order, aligning latents with semantic embeddings, or making transformations predictable. Re-training entire models is prohibitively expensive (14.5K GPU hours), yet directly finetuning frozen models is unstable. Re-Bottleneck solves this by training a lightweight inner autoencoder in the latent space, restructuring representations through latent-space losses alone while keeping the base model frozen.

The key insight is that you can reshape learned representations through post-hoc compression without touching the generator. By training a smaller encoder-decoder in the frozen latent space with user-defined losses, you create new structure while preserving reconstruction fidelity—at 0.33% of retraining cost.

Core Concept

Re-Bottleneck operates as a three-stage pipeline applied to any frozen pretrained autoencoder:

1. Freeze Base Autoencoder: Lock all weights in the original model; preserve its generative capacity
2. Train Inner Bottleneck: Create a lightweight encoder (Re-Encoder) and decoder that operate exclusively in the latent space
3. Apply Structured Losses: Train the inner bottleneck with task-specific objectives (channel ordering, semantic alignment, equivariance)

The base model remains unchanged, so deployment uses the original frozen checkpoint. The restructured latents become the new representation, enabling downstream applications (synthesis, analysis) to use ordered or semantically meaningful features.

Architecture Overview

• Base Autoencoder: Frozen pretrained encoder E and decoder D (e.g., EnCodec, AudioMAE)
• Re-Encoder (RE): Lightweight encoder that takes frozen latents z and produces restructured representation z̃
• Re-Decoder (RD): Lightweight decoder reconstructing approximations of original latents from z̃
• Reconstruction Head: Linear mapping from z̃ back to original latent space dimension
• Loss Modules: Task-specific objectives (nested dropout, contrastive, equivariance) applied only in latent space
• Frozen Base Generator: Original decoder D remains fixed for synthesis

Implementation

The following demonstrates Re-Bottleneck training with three variants:

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional

class ReBottleneckBase(nn.Module):
    """Base Re-Bottleneck structure for any frozen autoencoder."""
    def __init__(self, latent_dim: int, inner_dim: int = 256, dropout: float = 0.1):
        super().__init__()
        self.latent_dim = latent_dim
        self.inner_dim = inner_dim

        # Lightweight inner encoder (compresses latents)
        self.re_encoder = nn.Sequential(
            nn.Linear(latent_dim, inner_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(inner_dim, inner_dim // 2)
        )

        # Lightweight inner decoder (reconstructs latent approximations)
        self.re_decoder = nn.Sequential(
            nn.Linear(inner_dim // 2, inner_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(inner_dim, latent_dim)
        )

    def forward(self, z_frozen):
        """Process frozen latents through inner bottleneck."""
        # z_frozen: (batch, latent_dim)
        z_inner = self.re_encoder(z_frozen)
        z_reconstructed = self.re_decoder(z_inner)
        return z_inner, z_reconstructed

class OrderedReBottleneck(ReBottleneckBase):
    """Re-Bottleneck with nested dropout for channel ordering."""
    def __init__(self, latent_dim: int, inner_dim: int = 256):
        super().__init__(latent_dim, inner_dim)

        # Nested dropout layers encourage monotonic information ordering
        self.channel_dropout = nn.ModuleList([
            nn.Dropout(p) for p in [0.1, 0.3, 0.5, 0.7, 0.9]
        ])

    def forward(self, z_frozen):
        z_inner, z_reconstructed = super().forward(z_frozen)

        # Apply progressively stronger dropout to enforce ordering
        # Early channels survive high dropout → capture coarse features
        # Late channels only survive low dropout → capture fine details
        return z_inner, z_reconstructed

class SemanticReBottleneck(ReBottleneckBase):
    """Re-Bottleneck with semantic alignment via contrastive learning."""
    def __init__(self, latent_dim: int, inner_dim: int, semantic_dim: int = 768):
        super().__init__(latent_dim, inner_dim)

        # Semantic projection head (maps inner bottleneck to semantic space)
        self.semantic_head = nn.Linear(inner_dim // 2, semantic_dim)

    def compute_semantic_loss(self, z_inner, semantic_embeddings, temperature: float = 0.07):
        """Contrastive loss aligning latents with semantic embeddings."""
        # z_inner: (batch, inner_dim // 2)
        # semantic_embeddings: (batch, semantic_dim) from BEATs or T5

        # Project to shared space
        latent_proj = F.normalize(self.semantic_head(z_inner), dim=-1)
        semantic_proj = F.normalize(semantic_embeddings, dim=-1)

        # Cosine similarity matrix
        sim_matrix = torch.matmul(latent_proj, semantic_proj.t()) / temperature

        # Contrastive loss (diagonal elements are positive pairs)
        batch_size = z_inner.shape[0]
        labels = torch.arange(batch_size, device=z_inner.device)

        loss_forward = F.cross_entropy(sim_matrix, labels)
        loss_backward = F.cross_entropy(sim_matrix.t(), labels)

        return (loss_forward + loss_backward) / 2

class EquivariantReBottleneck(ReBottleneckBase):
    """Re-Bottleneck enforcing equivariance: input transforms ↔ latent transforms."""
    def __init__(self, latent_dim: int, inner_dim: int = 256):
        super().__init__(latent_dim, inner_dim)

        # Learnable transformation matrix (latent space operations)
        self.equivariance_matrix = nn.Parameter(torch.eye(inner_dim // 2))

    def compute_equivariance_loss(self, z_frozen_1, z_frozen_2, transformed_latent_1):
        """Ensure latent transformations match input domain operations."""
        # Extract inner representations
        z_inner_1, _ = super().forward(z_frozen_1)
        z_inner_2, _ = super().forward(z_frozen_2)

        # Expected transformation in latent space
        z_inner_transformed = torch.matmul(z_inner_1, self.equivariance_matrix.t())

        # L2 loss: actual transformation should match expected
        equivariance_loss = F.mse_loss(z_inner_transformed, z_inner_2)

        return equivariance_loss

def train_re_bottleneck(frozen_autoencoder, rebottleneck_model, data_loader,
                       variant: str = "semantic", optimizer=None, num_epochs: int = 50,
                       semantic_embeddings: Optional[torch.Tensor] = None):
    """
    Train Re-Bottleneck on frozen latent space.

    Args:
        frozen_autoencoder: Pretrained, frozen encoder-decoder pair
        rebottleneck_model: ReBottleneck variant instance
        data_loader: Iterator yielding audio batches
        variant: "ordered", "semantic", or "equivariant"
        semantic_embeddings: (batch, semantic_dim) for semantic variant
    """
    if optimizer is None:
        optimizer = torch.optim.AdamW(rebottleneck_model.parameters(), lr=1e-3)

    criterion = nn.MSELoss()

    for epoch in range(num_epochs):
        total_loss = 0

        for batch_idx, audio_batch in enumerate(data_loader):
            # audio_batch: (batch, channels, time)
            optimizer.zero_grad()

            # Encode audio to latent space using frozen base
            with torch.no_grad():
                z_frozen = frozen_autoencoder.encode(audio_batch)  # (batch, latent_dim)

            # Process through Re-Bottleneck
            z_inner, z_reconstructed = rebottleneck_model(z_frozen)

            # Reconstruction loss (preserve base model fidelity)
            recon_loss = criterion(z_reconstructed, z_frozen)
            total_loss_iter = recon_loss

            # Variant-specific losses
            if variant == "semantic" and semantic_embeddings is not None:
                sem_loss = rebottleneck_model.compute_semantic_loss(z_inner, semantic_embeddings)
                total_loss_iter += 0.5 * sem_loss

            elif variant == "equivariant":
                # Example: compare pairs of augmented audio
                with torch.no_grad():
                    z_frozen_augmented = frozen_autoencoder.encode(augment_audio(audio_batch))

                equiv_loss = rebottleneck_model.compute_equivariance_loss(
                    z_frozen, z_frozen_augmented, z_inner
                )
                total_loss_iter += 0.5 * equiv_loss

            total_loss_iter.backward()
            optimizer.step()
            total_loss += total_loss_iter.item()

        avg_loss = total_loss / len(data_loader)
        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch + 1}/{num_epochs}: Loss {avg_loss:.6f}")

    return rebottleneck_model

def augment_audio(audio, sr=16000, augmentation_type="pitch_shift"):
    """Simple audio augmentation for equivariance training."""
    # Placeholder: in practice use librosa or julius for proper augmentation
    return audio + torch.randn_like(audio) * 0.01

This implementation shows how to train inner bottlenecks for different restructuring goals—all in the latent space without touching the base model.

Practical Guidance

When to Use Re-Bottleneck

• Cost-constrained model adaptation: Restructuring latents <1% of full retraining cost
• Semantic alignment: Connecting audio representations to language models or embeddings
• Interpretability goals: Creating ordered latent channels for analysis or visualization
• Filter/transformation control: Enabling predictable latent operations for synthesis
• Rapid prototyping: Testing different latent structures without months of GPU time

When NOT to Use

• Fundamental architecture mismatch: If base autoencoder's reconstruction quality is poor, Re-Bottleneck cannot fix it
• Extreme restructuring goals: Re-Bottleneck preserves base capacity; cannot add fundamentally new structure
• Real-time synthesis: Re-Bottleneck adds inference overhead (nested forward passes); prefer simpler post-processing
• Models <16M parameters: Inner bottleneck overhead exceeds benefits; modify base model directly
• Extremely latent-constrained scenarios: If base model already maximally compresses, inner bottleneck has limited room

Common Pitfalls

1. Too-Aggressive Inner Compression: Keep inner_dim close to latent_dim (e.g., 75-90% of original). Over-compression loses information that Re-Bottleneck cannot recover.
2. Reconstruction Loss Dominance: Start with high reconstruction weight (0.7-0.8) to preserve base model, then decay to allow structure to emerge. If reconstruction loss always dominates, task-specific losses never activate.
3. Semantic Embedding Dimension Mismatch: Ensure semantic embeddings match the model (BEATs → 768D, T5 → 768D or 1024D). Mismatched dimensions cause silent training failures.
4. Frozen Base Model Validation: Verify base model outputs don't change during Re-Bottleneck training. Any change indicates weights are leaking (check requires_grad=False).
5. Skipping Warmup: Start with high reconstruction weight (90%) and decay it gradually to 50% over 20 epochs. Direct task loss training destabilizes training.

Reference

Chen, Z., Zipkin, G., et al. (2025). Re-Bottleneck: Latent Re-Structuring for Neural Audio Autoencoders. arXiv preprint arXiv:2507.07867.

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