Finding Archetypal Spaces Using Neural Networks

Abstract: Archetypal analysis is a data decomposition method that describes each observation in a dataset as a convex combination of "pure types" or archetypes. These archetypes represent extrema of a data space in which there is a trade-off between features, such as in biology where different combinations of traits provide optimal fitness for different environments. Existing methods for archetypal analysis work well when a linear relationship exists between the feature space and the archetypal space. However, such methods are not applicable to systems where the feature space is generated non-linearly from the combination of archetypes, such as in biological systems or image transformations. Here, we propose a reformulation of the problem such that the goal is to learn a non-linear transformation of the data into a latent archetypal space. To solve this problem, we introduce Archetypal Analysis network (AAnet), which is a deep neural network framework for learning and generating from a latent archetypal representation of data. We demonstrate state-of-the-art recovery of ground-truth archetypes in non-linear data domains, show AAnet can generate from data geometry rather than from data density, and use AAnet to identify biologically meaningful archetypes in single-cell gene expression data.

• The paper presents AAnet, a deep autoencoder framework that learns nonlinear archetypal spaces by embedding data within a convex simplex.
• It demonstrates superior performance over linear methods by accurately capturing geometric data structures in synthetic, image, and biological datasets.
• The approach enables uniform geometry-based generative modeling with high scalability, offering practical insights for diverse scientific applications.

Nonlinear Archetypal Analysis with Deep Networks: The AAnet Framework

Introduction and Motivation

Archetypal analysis (AA) is a geometric unsupervised learning method that decomposes data points as convex combinations of "archetypes," which represent the extrema or pure states of the data manifold. Conventional AA approaches—primarily linear, such as principal convex hull analysis (PCHA)—fit a simplex in the original feature space, representing each data point as a mixture of archetype vertices. However, when the data manifold is a nonlinear transformation of archetypal states (common in biological, physical, or image data), the linear simplex fit fails to correctly identify meaningful archetypes. Nonlinear methods, such as kernel archetypal analysis, rely on fixed transformations and lack adaptivity to the underlying data geometry.

This work proposes a novel problem formulation: learning an optimal nonlinear transformation that embeds data into a latent archetypal space explicitly bounded by a simplex, such that archetype extremality and convex representation are preserved and invertible. The solution, Archetypal Analysis network (AAnet), is a deep autoencoder with a specialized archetypal regularization on the latent layer, providing convex constraints without a fixed kernel or preprocessing transformation.

(Figure 1)

Figure 1: Schematic of AAnet showing nonlinear encoding of data into an archetypal simplex and its decoding for data analysis and generation.

AAnet supports inference of archetypes in nonlinearly embedded data, robust data generation that captures geometry beyond local density, and semantic exploration of mixtures for both visualization and downstream applications.

Problem Formulation and AAnet Architecture

Generalized Problem Statement

Let $\mathbf{X} \in \mathbb{R}^{n \times d}$ be the feature matrix. The goal is to learn a nonlinear transformation $f: \mathbb{R}^{d} \to \mathbb{R}^{k}$ (where $k$ is the number of archetypes) such that $f(\mathbf{x}_i)$ for each point $\mathbf{x}_i$ is a convex combination of simplex vertices, i.e., archetypes. $f$ should be approximately invertible to enable decoding back into the feature space. More precisely:

$$\min_{f, \{c_j\}} \sum_{i=1}^n \left\| f(\mathbf{x}_i) - \sum_{j=1}^k \alpha_{ij} c_j \right\|^2,$$

with $\sum_{j=1}^k \alpha_{ij} = 1$, $\alpha_{ij} \ge 0$, and $f$ invertible on $X$.

The archetypes are represented by one-hot vectors in the latent space, and the archetypal mixture of a point are its coordinates in this basis.

AAnet Neural Network

AAnet is a deep autoencoder with:

• Encoder: $E: \mathbb{R}^d \to \mathbb{R}^{k-1}$, computes the first $k-1$ convex coefficients;
• The $k$-th coefficient is determined implicitly: $\alpha_{ik} = 1 - \sum_{j=1}^{k-1} \alpha_{ij}$;
• Convexity constraints enforced as soft penalties during training: $\alpha_j \geq 0$ and $\sum_j \alpha_j \le 1$;
• Decoder: $D: \mathbb{R}^k \to \mathbb{R}^{d}$, reconstructs back to feature space.

The loss function is the sum of MSE reconstruction and convexity penalty terms, driving the latent layer to represent data inside a $k$-vertex simplex.

Optionally, Gaussian noise is injected in the latent space to control the "tightness" of the simplex to the data, analogous to robustness parameters in kernel AA methods.

Benchmarks: Nonlinear Data and Recovery of Archetypes

Synthetic Example: Triangle on a Sphere

A standard failure mode for linear AA occurs in the "triangle on a sphere" scenario, where data, originally a triangle in $\mathbb{R}^2$, is projected onto a nonlinear manifold (a sphere). Linear AA cannot identify the proper archetypes, as their convex hull in embedding space doesn't correspond to the true geometric extrema.

AAnet, by learning an adaptive nonlinear embedding, is able to consistently recover correct archetypes across increasing curvature levels. This is quantitatively validated by comparing MSE losses against ground-truth archetype positions and archetypal mixing coefficients, outperforming all competitor methods.

Figure 2: Comparison of AA methods in recovering nonlinear simplex structure—AAnet uniquely recovers ground truth archetypes as curvature increases.

Nonlinear Image Benchmarks: dSprites

Using the dSprites dataset, where image latent factors are affine (nonlinear in pixel space), AAnet is tested against other AA methods. Archetypal hearts, reconstructed archetypal spaces, and quantitative error metrics are evaluated. AAnet achieves the lowest archetypal space MSE across multiple data splits, with an 80% performance margin over the second-best method (“PCHA on AE”).

Figure 3: dSprites benchmarking—AAnet recovers correct archetypes with the lowest archetypal space error among all methods.

Data Generation and Geometry-Capturing

A critical conceptual advance is the decoupling of data geometry from data density in generative modeling. Standard VAEs and GANs generate samples according to observed density, often missing geometrically sparse regions.

AAnet, by sampling convex combinations in the archetypal latent space, generates data uniformly over the data manifold, regardless of original sampling density. This is quantitatively assessed with Maximum Mean Discrepancy (MMD); AAnet-generated samples yield a 56–64% reduction in discrepancy to ground-truth geometry versus GANs and VAEs.

Figure 4: Geometry-based generation—AAnet reconstructs manifold geometry for uniform generation, unlike GAN/VAE models.

Reproducibility, Scalability, and Practical Considerations

AAnet demonstrates high stability: repeated runs with different initializations recover consistent archetypes, outperforming random baselines by precise $R^2$ measures.

Latent Gaussian noise controls archetype “tightness” to data, permitting a trade-off between extremality and representation of archetypes.

AAnet is highly scalable. Empirically, its runtime scales sub-linearly with dataset size, remaining tractable for samples up to hundreds of thousands of data points, whereas classical AA methods exhibit exponential scaling beyond 50,000 samples.

Visualization of the archetypal space is achieved through a specialized, efficient MDS interpolation method, which is orders of magnitude faster than standard MDS on large datasets, allowing high-throughput exploratory analysis.

Figure 5: Efficient MDS-based visualization for archetypal structure scales to tens of thousands of data points.

Applications to Biological Data: Single-Cell and Microbiome Profiling

In single-cell RNA-seq analysis of tumor-infiltrating lymphocytes, AAnet identifies functionally relevant immune archetypes, each corresponding to specific gene expression and phenotype modules. Distance from a cell to archetype vertices in latent space correlates with biological gradients (e.g., exhaustion, naïvety, cytotoxicity), validated through canonical markers and signature genes.

Figure 6: Application to single-cell lymphocyte data—archetypes correspond to distinct T-cell states and gene expression gradients.

AAnet also captures the continuous spectrum of gut microbiome diversity, uncovering both previously defined enterotypes and novel composition regimes—each associated with specific bacterial taxa enrichment—along gradients described by latent archetypal mixtures. This expands the classic enterotype paradigm to a higher-dimensional, nonlinear landscape.

Figure 7: AAnet recapitulates microbiome compositional gradients, mapping individuals along non-discrete archetypal axes.

Model Selection and Hyperparameterization

The optimal number of archetypes can be inferred directly from loss curve knee-point analysis, with the plateau behavior matching the true underlying generative process dimensionality. Latent noise $\sigma$ is used to regulate archetype extremality.

Network parameters are standard for deep autoencoder architectures, with LeakyReLU activations and Adam optimization. All archetypal constraints are implemented as differentiable soft penalties, enabling efficient GPU-based training.

Theoretical and Practical Implications

By explicitly learning a non-linear transformation into an archetypal simplex, AAnet generalizes and subsumes prior AA paradigms. Its architectural design combines rigorous geometric constraints with the flexibility and scalability of deep networks.

Practically, AAnet enables:

• Accurate recovery of underlying mechanisms in systems where traits arise as nonlinear combinations of pure states (e.g., cell differentiation, ecological profiles, market segmentation).
• Uniform, geometry-based generative modeling for simulation, imputation, and perturbation studies, robust to density artifacts and sampling biases.

Theoretically, this establishes a practical path for nonlinear extreme point analysis, potentially extending to supervised and semi-supervised tasks, as well as integration with graph and manifold learning frameworks. Further, the decoupling of geometry from density may inspire generative approaches in computer vision, structure learning, and causal inference.

Conclusion

"Finding Archetypal Spaces Using Neural Networks" establishes a deep learning solution for nonlinear archetypal analysis that is uniquely robust to data nonlinearity, efficient on large-scale datasets, and provides interpretable, geometry-driven generative modeling. The work sets a structural direction for future research on interpretable, geometry-preserving unsupervised learning in high-dimensional data, particularly in domains where understanding pure states and their mixtures underlies scientific discovery.