Rare-Aware Autoencoding: Reconstructing Spatially Imbalanced Data

Abstract

Autoencoders can be challenged by spatially non-uniform sampling of image content. This is common in medical imaging, biology, and physics, where informative patterns occur rarely at specific image coordinates, as background dominates these locations in most samples, biasing reconstructions toward the majority appearance. In practice, autoencoders are biased toward dominant patterns resulting in the loss of fine-grained detail and causing blurred reconstructions for rare spatial inputs especially under spatial data imbalance. We address spatial imbalance by two complementary components: (i) self-entropy-based loss that upweights statistically uncommon spatial locations and (ii) Sample Propagation, a replay mechanism that selectively re-exposes the model to hard to reconstruct samples across batches during training. We benchmark existing data balancing strategies, originally developed for supervised classification, in the unsupervised reconstruction setting. Drawing on the limitations of these approaches, our method specifically targets spatial imbalance by encouraging models to focus on statistically rare locations, improving reconstruction consistency compared to existing baselines. We validate in a simulated dataset with controlled spatial imbalance conditions, and in three, uncontrolled, diverse real-world datasets spanning physical, biological, and astronomical domains. Our approach outperforms baselines on various reconstruction metrics, particularly under spatial imbalance distributions. These results highlight the importance of data representation in a batch and emphasize rare samples in unsupervised image reconstruction. We will make all code and related data available.

Figures (10)

• Figure 1: Spatial data imbalance.Top: example frames from a pendulum dataset at different angles. Due to damping, the pendulum spends more time near equilibrium, making some spatial configurations much more frequent than others. Top-right: per-pixel variance over the full dataset, highlighting regions with high motion-induced variability. Bottom-left: temporal intensity traces for four representative pixels: a static background pixel (blue), an extreme-angle pixel (red), a mid-position pixel (green), and an equilibrium pixel (orange). Variance increases as pixels approach the equilibrium region, since the pendulum passes there most often. Bottom-right: histogram of pixel intensities (for the selected pixels), showing the highly non-uniform occupancy of the state space and the resulting spatial imbalance.
• Figure 2: Addressing reconstruction under spatial data imbalance. Consider a dataset of simulated pendulum dynamics. Black squares are input or reconstructed frames. Left-top: Input distributions exhibit a strong bias, with common, overrepresented samples (blue) and underrepresented ones (red). Left-bottom: Standard autoencoder (AE) backbone (encoder–latent–decoder). The SOTA (cyan) approach tends to converge toward smooth reconstructions dominated by frequent samples. Right: Our approach (orange) introduces two complementary mechanisms: (i) a self-entropy loss that upweights high-surprisal pixels across the batch within each image, and (ii) sample propagation that replays high-loss (hard) samples across epochs, ensuring rare events are repeatedly observed. The combined effect yields sharper and more balanced reconstructions.
• Figure 3: Sample propagation algorithm. First, for each input sample, the model produces a reconstruction error. The color scale encodes error magnitude: red denotes higher error, while decreasing progressively to cooler hues. Second, we compute a error per sample and sort/rank the batch accordingly, yielding an ordered list from highest to lowest reconstruction error. Finally, a propagation step selects the $M$ top-ranked (hardest) samples for targeted optimization and concatenates them to the next batch; gradients are then applied (loss backward) to update the network with an emphasis on hard examples.
• Figure 4: Comparison between balanced and imbalanced reconstruction scenarios. The MNIST dataset (left) represents a balanced case where samples are more uniformly distributed across the spatial domain. In contrast, the damped pendulum dataset (right) introduces imbalance, as certain image spatial locations are severely underrepresented across the dataset and within training batches. Our proposed spp+entropy framework reconstructs both balanced and imbalanced data faithfully, while the standard MLP AE pipeline performs well only on homogeneous datasets such as MNIST but fails to preserve structure under imbalance.
• Figure 5: (Left) Qualitative comparison. The first row shows the ground truth (GT). Methods are arranged by rows: (Left) comparison of baselines and our entropy-based methods (entropy*, entropy+spp*); (right) baselines and our underrepresented-sample-focusing spp variants (spp2*, spp4*, spp8*). Columns include both an Underrepresented sample (Us) and an Average sample (As) to visualize behavior on rare versus frequent cases. The layout highlights how different objectives affect sharp structures and textures: the left column makes visible the tendency of some losses to approach a dataset-average appearance, whereas the spp block illustrates how prioritizing underrepresented samples influences detail retention and stability across memory sizes $k$.