Hierarchical Conditioning of Diffusion Models Using Tree-of-Life for Studying Species Evolution (2408.00160v1)

Abstract: A central problem in biology is to understand how organisms evolve and adapt to their environment by acquiring variations in the observable characteristics or traits of species across the tree of life. With the growing availability of large-scale image repositories in biology and recent advances in generative modeling, there is an opportunity to accelerate the discovery of evolutionary traits automatically from images. Toward this goal, we introduce Phylo-Diffusion, a novel framework for conditioning diffusion models with phylogenetic knowledge represented in the form of HIERarchical Embeddings (HIER-Embeds). We also propose two new experiments for perturbing the embedding space of Phylo-Diffusion: trait masking and trait swapping, inspired by counterpart experiments of gene knockout and gene editing/swapping. Our work represents a novel methodological advance in generative modeling to structure the embedding space of diffusion models using tree-based knowledge. Our work also opens a new chapter of research in evolutionary biology by using generative models to visualize evolutionary changes directly from images. We empirically demonstrate the usefulness of Phylo-Diffusion in capturing meaningful trait variations for fishes and birds, revealing novel insights about the biological mechanisms of their evolution.

• The paper introduces Phylo-Diffusion, leveraging hierarchical embeddings to condition latent diffusion models for visualizing evolutionary traits.
• The methodology employs multi-level phylogenetic embeddings to capture species’ inherited traits and guide generative image synthesis.
• Experiments with trait masking and swapping demonstrate the framework’s potential in revealing key evolutionary insights and validating phylogenetic relationships.

Hierarchical Conditioning of Diffusion Models Using Tree-of-Life for Studying Species Evolution

Introduction

The paper "Hierarchical Conditioning of Diffusion Models Using Tree-of-Life for Studying Species Evolution" by Khurana et al. proposes Phylo-Diffusion, a novel framework integrating phylogenetic knowledge into diffusion models to paper evolutionary traits in various species. Utilizing the Tree of Life, the authors introduce hierarchical embeddings (HIER-Embeds) to condition latent diffusion models, allowing for the visualization and analysis of trait variations across species' evolutionary timelines. The paper's contributions span methodological advancements in generative modeling, empirical evaluations, and new insights into evolutionary biology.

Methodology

The crux of Phylo-Diffusion lies in its hierarchical embedding strategy, HIER-Embed, which structures the embedding space of diffusion models using multi-level representations of phylogenetic knowledge. Each species' phylogenetic information is encoded as a sequence of vectors across four ancestry levels. This approach helps in capturing trait information inherited at different evolutionary periods.

Hierarchical Embedding

Phylo-Diffusion employs HIER-Embeds to condition the latent diffusion models. Each level embedding captures different evolutionary stages, ensuring a species' hierarchical trait representations are preserved. Combining these embeddings offers a unified representation that conditions the diffusion models to generate synthetic images reflective of a species' evolutionary traits.

Generative Modeling

The latent diffusion model (LDM) framework is utilized for generating high-quality images, operating in compressed latent spaces, which accelerates training and inference times. The model’s conditioning mechanism incorporates the hierarchical embeddings through cross-attention mechanisms, significantly enhancing its ability to generate biologically meaningful images conditioned on phylogenetic data.

Experiments

The authors propose two novel experiments, trait masking and trait swapping, to perturb the embedding space and analyze resulting changes in generated images' traits.

1. Trait Masking: This involves substituting level-specific embeddings with noise to paper trait erasures reflective of genetic knockouts. The model is evaluated for how well it retains common traits up to the masked level.
2. Trait Swapping: Inspired by gene editing, this experiment swaps embeddings at a specific level with sibling species' embeddings. The resulting images highlight trait differences that arose due to evolutionary branching at that level.

Results

The effectiveness of Phylo-Diffusion is demonstrated using datasets of fish and bird images. Key findings include:

• Image Quality and Classification: Phylo-Diffusion generates images with FID scores comparable to state-of-the-art models. Additionally, synthetic images are classified with high F1 scores, indicating that the generated images faithfully represent the intended species' traits.
• Probability Distributions: Trait masking experiments reveal significant changes in classification probabilities, accurately reflecting the phylogenetic relationships. Species within the same subtree exhibit more pronounced probability increases compared to those outside, validating the hierarchical structuring in embeddings.
• Trait Insights: Trait swapping experiments yield profound insights into evolutionary traits. For instance, distinguishing features like the absence of barbels and changes in fin structures are consistently observed in swapped images, providing visual hypotheses for phylogenetic studies.

Implications and Future Work

The methodological advances proposed in this paper hold significant implications for both AI and evolutionary biology:

• Practical AI Applications: The hierarchical conditioning mechanism can be extended to other domains requiring phylogenetic or hierarchical data representations, such as oncology or genomics.
• Evolutionary Biology: Phylo-Diffusion enables rapid hypothesis generation and testing by visualizing evolutionary changes, which traditionally requires labor-intensive and time-consuming empirical studies. This can accelerate discoveries in systematics and phenotypic trait analysis.

Future research directions include refining the approach to handle trees with varying levels of discretization, addressing convergent evolutionary changes, and incorporating uncertainty estimates in ancestral state reconstructions.

Conclusion

Phylo-Diffusion represents a significant step toward integrating generative models with structured biological knowledge. By leveraging hierarchical embeddings, this framework provides an innovative lens for visualizing and understanding species evolution directly from images. The proposed trait masking and trait swapping experiments offer new methodologies for studying evolutionary biology, facilitating more nuanced insights into the complexity of species' evolutionary histories.