VAE Full-Band Upsampling

• VAE-based full-band upsampling is a technique that combines spatial and spectral loss functions to accurately reconstruct both low- and high-frequency signal details.
• The method employs innovative decoder architectures, including transposed convolution and nearest-neighbor interpolation, to enhance high-frequency fidelity.
• It integrates latent-domain upsampling for audio processing, achieving significant computational efficiency while maintaining reconstruction quality.

Variational Autoencoder (VAE)-based full-band upsampling refers to a family of techniques that leverage the latent bottleneck and generative capacity of VAEs to reconstruct signals—typically images or audio—at higher bandwidth or resolution, with explicit emphasis on matching both low- and high-frequency information. These methods incorporate architectural modifications and loss functions designed to address deficiencies in conventional VAE reconstructions, especially the tendency to lose or blur fine spectral details. VAE-based full-band upsampling can be implemented either in the output (data) domain, where frequency-domain regularization is enforced directly, or entirely within the latent domain, where all upsampling operations occur on a compressed representation prior to decoding.

1. Mathematical Characterization of Spectral Regularization

A central challenge in VAE-based upsampling is the accurate reproduction of both spatial and frequency components. The primary method for spectral fidelity is the combination of spatial and spectral-domain loss functions. For an image $x \in \mathbb{R}^{H \times W \times C}$ and its VAE reconstruction $\hat{x}$, a combined reconstruction loss is defined as: $$L(x, \hat{x}) = \alpha L_s(x, \hat{x}) + (1-\alpha) L_f(x, \hat{x}), \quad \alpha \in [0, 1]$$ where:

• $L_s(x, \hat{x}) = \text{BCE}(x, \hat{x})$ is the pixel-wise binary cross-entropy (spatial loss),
• $L_f(x, \hat{x})$ is the mean squared error over the real and imaginary parts of the 2D discrete Fourier transforms of $x$ and $\hat{x}$, applied channel-wise:

$$L_f(x, \hat{x}) = \frac{1}{n} \sum_{u,v,c} | \operatorname{Re}(\mathcal{F}_c[x]_{u,v}) - \operatorname{Re}(\mathcal{F}_c[\hat{x}]_{u,v}) |^2 + \frac{1}{n} \sum_{u,v,c} | \operatorname{Im}(\mathcal{F}_c[x]_{u,v}) - \operatorname{Im}(\mathcal{F}_c[\hat{x}]_{u,v}) |^2$$

where $n = H \times W \times C$.

The parameter $\alpha$ mediates the trade-off between spatial detail and spectral (frequency) fidelity, with typical empirically tuned values around $\alpha \approx 0.5$ yielding effective results. This joint loss enforces not only power matching but also phase accuracy for each 2D frequency component, imposing penalties for both low- and high-frequency mismatches (Björk et al., 2022).

2. Decoder Upsampling Architectures

Full-band VAE upsampling performance depends critically on the upsampling architecture within the decoder. Two schemes are contrasted:

• Transposed convolution (‘deconvolution’): Each upsampling block employs a ConvTranspose2d layer (e.g., kernel size 4, stride 2, padding 1), followed by batch normalization and ReLU, except the final sigmoid output. This architecture propagates learned upsampling throughout all layers.
• Nearest-neighbor interpolation plus convolution (‘N.1.5’): All but the final block use transposed convolutions; the last upsampling replaces ConvTranspose2d with (i) 2× nearest-neighbor interpolation (parameter-free), (ii) a 5×5 convolution + batch norm + ReLU, and (iii) a final 3×3 convolution projecting to output channels, concluded with a sigmoid. Bilinear upsampling was not tested, but is frequently cited as a common alternative.

Empirical findings indicate that substituting only the final upsampling block with “N.1.5” often boosts high-frequency fidelity but does not alone guarantee consistent improvement across all spectral domains. This suggests the choice of upsampling scheme is necessary but not sufficient for full-band spectral accuracy (Björk et al., 2022).

3. VAE Objective Integration and Training Paradigms

The overall VAE loss combines the hybrid spatial+spectral reconstruction objective with the standard latent regularization: $$L_{\text{total}} = \mathbb{E}_{x \sim \text{data}} \left[ \mathbb{E}_{z \sim q_\phi(z|x)} \big[ L(x, \text{decoder}(z)) \big] + \beta\, \text{KL}(q_\phi(z|x) \| p(z)) \right]$$ Here, $\beta$ is typically set to 1.0 as in standard VAEs. The weighting parameter $\alpha$ is selected on a validation set but generally $\alpha \approx 0.5$ provides robust trade-offs. No additional loss weighting is introduced beyond $\alpha$.

For audio, “Learning to Upsample and Upmix Audio in the Latent Domain” (Bralios et al., 31 May 2025) demonstrates that all upsampling may instead be performed directly in the latent space. Here, the upsampler $F_\theta$ maps low-band latent representations $\mathbf{z}_\text{low}$ to full-band $$\hat{\mathbf{z}}_\text{high}=F_\theta(\mathbf{z}_\text{low}, \mathbf{c})$$. Training occurs entirely in latent space, using an $L_1$ reconstruction term, an optional adversarial loss, and—where applicable—a variational KL with conditioning encoder.

4. Experimental Evaluation Protocols

VAE-based full-band upsampling is benchmarked on both synthetic and real datasets:

• Image domain (Björk et al., 2022):
  • Datasets: SHAPE (synthetic, 32×32 grayscale geometric shapes), MNIST (32×32), CelebA (64×64, RGB, center-cropped).
  • Evaluation metrics: Spatial RMSE, 2D FFT-domain RMSE, and 1D azimuthal-integration (AI) power spectrum RMSE.
  • Objectives: Vanilla-VAE (BCE only), Watson-DFT loss, AI loss, and 2D FFT loss.
  • Ablations: All loss functions paired with both transposed convolution and “N.1.5” upsampling.
• Audio domain (Bralios et al., 31 May 2025):
  • Datasets: 10,000 hours of instrumental music, 44.1 kHz.
  • Tasks: Bandwidth extension (BWE) from 22.05 to 44.1 kHz; mono-to-stereo (M2S) upmixing.
  • Metrics: STFT-distance (“STFT-D”), mel-distance, measured on FMA-small. Computational costs quantified in GFLOPS per second of audio.
  • Baselines: Raw-audio upsamplers (Aero, MusicHiFi) versus Re-Encoder (latent upsampler).

Domain	Data type	Key metrics	Main baseline(s)	Upsampling method
Image	2D images	Spatial RMSE, 2D FFT RMSE, 1D AI RMSE	Vanilla-VAE, SR losses	Transposed conv, N.1.5
Audio	Waveform audio	STFT-D, mel-distance, FLOPS	Aero, MusicHiFi	Latent ConvNeXt, GAN

5. Empirical Observations and Comparative Results

For images, spectral regularization with the 2D FFT loss consistently outperforms both vanilla VAE and specialized perceptually-motivated losses (Watson-DFT, AI loss) across almost all datasets and evaluation domains, particularly in reconstruction of high-frequency detail. When “N.1.5” is used for only the final upsampling, modest further improvements in high-frequency RMSE and sharper spectral alignment are observed, though the effect is less uniform than with spectral loss (Björk et al., 2022).

Qualitative inspection reveals that FFT-regularized VAE reconstructions exhibit sharper edges and a power spectrum closely matching that of the ground truth, while vanilla VAE results are comparatively blurred with underrepresented high-frequency energy.

In audio, latent-domain upsampling (Re-Encoder) offers matched or improved STFT/distortion metrics over raw-audio upsampling, at up to 200× lower computational cost (0.4–1.6 GFLOPS for Re-Encoder versus 85–111 GFLOPS for baseline methods). The paradigm is empirically validated for both bandwidth extension and mono-to-stereo upmixing, with perceptual quality as measured by STFT-D closely approaching raw-audio optimized methods (Bralios et al., 31 May 2025).

6. Methodological Guidelines and Limitations

Practitioners are advised:

• Integrate a 2D FFT-based spectral MSE into the reconstruction loss, with $\alpha \approx 0.5$, for robust full-band fidelity.
• Use “N.1.5” upsampling for improved high-frequency reconstruction, but do not rely exclusively on upsampling strategy for spectral accuracy.
• Always evaluate models in spatial, 2D FFT, and 1D AI domains, as improvements do not transfer uniformly.
• In the latent domain (audio), perform all upsampling operations before decoding to achieve significant efficiency gains, bearing in mind that the final output fidelity remains bounded by the base autoencoder’s capacity.
• Be aware that deterministic latent-domain models may tend to generate mean outputs lacking detail unless adversarial or variational objectives are included.

Limitations include the irreducible reconstruction error set by the base (frozen) codec and the possible over-smoothing in the absence of richer generator/discriminator interactions. Future research directions include extending latent-domain upsampling to additional tasks (e.g., denoising, source separation) and exploring adaptive fine-tuning of decoders.

7. Impact and Future Considerations

The introduction of spectral-domain regularization and latent-domain processing establishes a paradigm for VAE-based full-band upsampling that is both computationally efficient and spectrally accurate. By combining a straightforward 2D FFT-based loss with careful upsampling design, or by operating fully in latent space, modern VAE systems can approach or exceed state-of-the-art performance in both image and audio modalities, with substantially reduced engineering complexity and resource demand (Björk et al., 2022, Bralios et al., 31 May 2025). This suggests broad applicability across generative modeling tasks where high-fidelity reconstructions and efficiency are paramount, and points to new directions in codec-agnostic neural upsampling and latent generative processing.