Machine Learning Engineer Intern MS PhD, Summer

Position: Machine Learning Engineer Intern (MS/PhD, 2027 Summer) Train and fine-tune a large multi-task transformer over driving-log sequences, in JAX/Flax on TPUs - iterating on fine-tuning strategies, training data mixtures, and losses to improve evaluation quality for the hillclimbing workflow Design and run rigorous offline and end-to-end evaluations - PR-AUC, calibration quality, and metric sensitivity on real hillclimbing A/B runs - and build the dataset and evaluation pipelines needed to produce them Land production-quality code in a shared, high-traffic codebase, and communicate results through a design doc, team deep dives, and a final intern presentation, partnering with UEM Core, Data Science, and release-eval stakeholders

Benefits Medical, dental, and vision insurance for employees and dependents Employee assistance programs focused on mental health Personalized workplace adjustments for diverse needs and abilities, including physical, mental, and neurodivergent considerations Access to mental health apps Onsite wellness centers Medical advocacy program for transgender employees Second medical opinion for you and your loved ones Counseling services Support programs including menopause benefit Competitive compensation Regular bonus and equity performance grant opportunities Generous 401(k) and regional retirement plans 1-on-1 financial coaching Annual cross-company compensation review and pay equity analysis Fertility and growing family assistance Parental leave and baby bonding leave Backup childcare Elder care and support Survivor income benefit Caregiver leave Paid time off, including vacation, bereavement, sick leave, parental leave, disability, and holidays Jury duty leave Military leave Hybrid work model with remote work opportunities also available Educational reimbursement Peer learning and coaching platform Donation matching programs Employee resource groups Volunteer hours Internal community groups and local culture clubs Inspiring spaces to work, recharge, and collaborate with fellow Waymonauts On-site meals and snacks Fitness centers, massage programs, and ergonomic support On-demand fitness, wellbeing, and cooking classes Commuter benefits Hands-on experience training and evaluating deep learning models in a modern framework (JAX, PyTorch or Tensor Flow), including building data pipelines, choosing losses, and debugging training runs Strong programming skills in C++/Python, plus a solid grounding in ML fundamentals: precision/recall trade-offs, class imbalance, evaluation metric selection, and rigorous experiment design Currently enrolled in an PhD or MS program in Computer Science, Machine Learning or a related field, returning to the program after the internship Authorship of published papers in top-tier AI/ML, data mining, or computer vision conferences (e.g., NeurIPS, ICML, ICLR, KDD, CVPR, CoRL, SIGMOD, VLDB, ACL) Research or applied experience with transformer and sequence models, multi-task learning, transfer learning or domain adaptation, and parameter-efficient fine-tuning of large pretrained models

Experience with JAX/Flax, distributed training on TPUs or GPUs, and large-scale data processing (Map Reduce-style pipelines, SQL) for building training and evaluation datasets Familiarity with autonomous driving, robotics, or simulation; and/or with probability calibration, uncertainty quantification, importance sampling, active learning, or rare-event and imbalanced-data modeling