Head of AI Engineering - Lead ML, Knowledge Graphs & Scale

highlightai; San Francisco (CA); On-site; USD 220,000 - 340,000; Full time

Benefits offered by this job

Salary + equity

Health benefits

Flexible PTO

Team meals

Relocation

Job summary

Highlight is building a shared intelligence layer for the modern workforce. We seek a Head of AI Engineering to own the architecture, quality, and roadmap for how we ingest data, model, and surface context—from data pipelines to the context graph that makes our system useful and scalable.

This is a hands-on leadership role: you’ll write code, set technical direction, mentor engineers, and partner with Backend and Product teams to ship features that feel like magic.

Qualifications

• 7+ years in ML/AI engineering with 3+ years leading teams or major technical initiatives
• Master's or PhD in Computer Science, Machine Learning, Statistics, Applied Math or equivalent experience
• Published research (or production ML) in knowledge graphs, memory systems, retrieval, or continual learning
• Experience at top AI labs (e.g., DeepMind, OpenAI, Anthropic, Meta FAIR, xAI) or leading AI companies (Google, Cursor, Perplexity)
• Strong architectural chops: designed and scaled ML infrastructure from scratch
• Hands-on proficiency with PyTorch, TensorFlow, RAG, vector DBs, cloud infra
• A player-coach mentality: mentoring engineers and coding
• Ownership mentality and in-office SF base
• Excellent communication and collaboration in lean teams
• Dedication to high-quality product shipping

Responsibilities

• Own the context graph architecture and memory systems end-to-end
• Research, prototype and deploy state-of-the-art ML models (transformers, CV, time series, RL)
• Lead full ML lifecycle: data pipelines, training, evals, deployment, monitoring, scalability, latency, retraining, A/B testing
• Own output quality across Highlight chat, actions, and workflows
• Build and lead a world-class ML Engineering team; establish ML practices
• Stay current on knowledge representation, retrieval, and continual learning; publish when appropriate

Skills

Leadership

Hands-on coding

Communication

Ownership

Collaboration

Education

Master's or PhD in CS/ML or related

Publications or production ML experience

Tools

PyTorch

TensorFlow

RAG

Vector DBs

Cloud infra

Job description

Highlight is building a shared intelligence layer for the modern workforce. We seek a Head of AI Engineering to own the architecture, quality, and roadmap for how we ingest data, model, and surface context—from data pipelines to the context graph that makes our system useful and scalable.

This is a hands-on leadership role: you’ll write code, set technical direction, mentor engineers, and partner with Backend and Product teams to ship features that feel like magic.