re-bottleneck-latent-restructuring

name	re-bottleneck-latent-restructuring
title	Re-Bottleneck: Latent Re-Structuring for Neural Audio Autoencoders
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2507.07867
keywords	["Audio Autoencoders","Latent Space","Post-hoc Modification","Semantic Alignment"]
description	Restructure latent representations in pretrained audio autoencoders without full retraining. Apply three variants—ordered, semantic, and equivariant—to enforce structure like channel ordering, semantic alignment, or filter correspondence. Achieves 20-60% semantic gains in under 48 GPU hours versus 14.5K hours for full retraining.

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
 torch.nn  nn
 torch.nn.functional  F
 typing  

 (nn.Module):
    
     ():
        ().__init__()
        .latent_dim = latent_dim
        .inner_dim = inner_dim

        
        .re_encoder = nn.Sequential(
            nn.Linear(latent_dim, inner_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(inner_dim, inner_dim // )
        )

        
        .re_decoder = nn.Sequential(
            nn.Linear(inner_dim // , inner_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(inner_dim, latent_dim)
        )

     ():
        
        
        z_inner = .re_encoder(z_frozen)
        z_reconstructed = .re_decoder(z_inner)
         z_inner, z_reconstructed

 ():
    
     ():
        ().__init__(latent_dim, inner_dim)

        
        .channel_dropout = nn.ModuleList([
            nn.Dropout(p)  p  [, , , , ]
        ])

     ():
        z_inner, z_reconstructed = ().forward(z_frozen)

        
        
        
         z_inner, z_reconstructed

 ():
    
     ():
        ().__init__(latent_dim, inner_dim)

        
        .semantic_head = nn.Linear(inner_dim // , semantic_dim)

     ():
        
        
        

        
        latent_proj = F.normalize(.semantic_head(z_inner), dim=-)
        semantic_proj = F.normalize(semantic_embeddings, dim=-)

        
        sim_matrix = torch.matmul(latent_proj, semantic_proj.t()) / temperature

        
        batch_size = z_inner.shape[]
        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)

         (loss_forward + loss_backward) / 

 ():
    
     ():
        ().__init__(latent_dim, inner_dim)

        
        .equivariance_matrix = nn.Parameter(torch.eye(inner_dim // ))

     ():
        
        
        z_inner_1, _ = ().forward(z_frozen_1)
        z_inner_2, _ = ().forward(z_frozen_2)

        
        z_inner_transformed = torch.matmul(z_inner_1, .equivariance_matrix.t())

        
        equivariance_loss = F.mse_loss(z_inner_transformed, z_inner_2)

         equivariance_loss

 ():
    
     optimizer  :
        optimizer = torch.optim.AdamW(rebottleneck_model.parameters(), lr=)

    criterion = nn.MSELoss()

     epoch  (num_epochs):
        total_loss = 

         batch_idx, audio_batch  (data_loader):
            
            optimizer.zero_grad()

            
             torch.no_grad():
                z_frozen = frozen_autoencoder.encode(audio_batch)  

            
            z_inner, z_reconstructed = rebottleneck_model(z_frozen)

            
            recon_loss = criterion(z_reconstructed, z_frozen)
            total_loss_iter = recon_loss

            
             variant ==   semantic_embeddings   :
                sem_loss = rebottleneck_model.compute_semantic_loss(z_inner, semantic_embeddings)
                total_loss_iter +=  * sem_loss

             variant == :
                
                 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 +=  * equiv_loss

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

        avg_loss = total_loss / (data_loader)
         (epoch + ) %  == :
            ()

     rebottleneck_model

 ():
    
    
     audio + torch.randn_like(audio) *