Senior/Staff Machine Learning Researcher

Terra AI; Palo Alto (CA); On-site; USD 100,000 - 160,000; Full time

Job summary

Terra AI is seeking a talented deep learning engineer or scientist in Palo Alto, California to lead the development of generative models for 3D geological structures and decision-making in clean energy applications.

The ideal candidate will have extensive experience with PyTorch, a strong foundation in developing large deep learning architectures, and skills in data curation. Responsibilities include designing and iterating on diffusion models, working with geophysical data, and collaborating with project teams.

Qualifications

• Extensive experience with PyTorch, including custom module creation.
• Ability to design and implement large-scale deep learning models.
• Hands-on skills in creating and maintaining high-quality datasets.
• Proficient in software development best practices.

Responsibilities

• Design and iterate diffusion models for 3D geological models.
• Condition generation on geophysical data and other observations.
• Inform synthetic data generation for model performance.
• Collaborate with project teams for specific real-world projects.

Skills

PyTorch

Deep Learning

Data Curation

Software Engineering

Job description

Role description

In the same way image generators have shown the remarkable ability to produce a diverse set of realistic pictures conditioned on a text prompt (and other inputs), we are developing a generative model that produces 3D geological models conditioned on geophysical surveys, borehole measurements, and other forms of physical observation. The outputs of the generative model capture what we know and don’t know about the state of the subsurface, allowing explorers to make maximally informed decisions about how and where to explore for critical resources.

We are looking for a talented deep learning engineer or scientist to lead the development of this model that will revolutionize decision-making in the earth subsurface for a wide range of clean energy applications.

Role Responsibilities

• Design, train, test, and iterate on diffusion models for 3D geological models
• Design, train, test, and iterate on an approach for conditioning generation on geophysical data and other observations
• Inform the generation of synthetic data to improve model performance
• Adapt diffusion modeling approach to specific real-world projects in collaboration with project teams.

Qualifications

Required Qualifications:

• Extensive PyTorch Experience
  • Deep understanding of PyTorch, including writing custom modules, optimizing training, and debugging issues in large-scale models.
• Expertise in Developing Large Deep Learning Models from Scratch
  • Proven ability to design, implement, and train complex deep learning architectures from the ground up.
• Data Curation Skills
  • Hands‑on experience in creating, cleaning, and maintaining high-quality datasets tailored for machine learning applications.
• Strong Software Engineering and Design Experience
  • Proficient in software development best practices, including version control, testing, and code optimization.
  • Familiarity with designing scalable and maintainable systems.

Nice-to-haves

• Experience with Generative Models
  • Familiarity with generative architectures, particularly diffusion models, and an emphasis on posterior sampling methods.
• Knowledge of Transformer Architectures
  • Experience building and training transformers, especially in applications involving 3D data.
• Scaling Models Across Large GPU Clusters
  • Expertise in parallelizing models across multiple GPUs and optimizing distributed training pipelines.
• Cloud Infrastructure Expertise
  • Experience setting up, managing, and optimizing cloud environments for machine learning workloads, including provisioning resources and managing costs.