Amir Ziashahabi

PhD Student, Electrical Engineering
University of Southern California

I work on , mostly on making them faster and cheaper to run.

Portrait of Amir Ziashahabi

About

I am a PhD student at the University of Southern California, advised by Prof. Salman Avestimehr. My research focuses on making machine learning systems more efficient, secure, and accessible. I work on improving large language model (LLM) inference through techniques such as speculative decoding and developing methods for fine-tuning on edge devices. I have also worked on privacy-preserving machine learning and building distributed systems for federated learning. I am passionate about bridging the gap between cutting-edge ML research and practical, deployable systems.

Interests: Efficient Machine Learning, Privacy-Preserving Machine Learning, Large Language Models, Distributed Systems.

Publications

* equal contribution
  1. ASPIRE: Asynchronous Batched Self-Speculative Decoding for Long-Context LLM Inference

    Amir Ziashahabi*, Hossein Entezari Zarch*, Lei Gao, Murali Annavaram, Salman Avestimehr

    COLM 2026PaperCode
  2. Spend Bits Where Queries Look: KV Cache Vector Quantization with Attention-Preserving Transforms

    Samuel Fernández-Menduiña*, Amir Ziashahabi*, Eduardo Pavez, Antonio Ortega, Salman Avestimehr

    Under reviewPaperCode
  3. ORBIT: Training-Free Multi-Attribute Behavioral Steering via Orthogonal Subspace Rotation

    Narges Ghasemi, Amir Ziashahabi, Salman Avestimehr, Jonathan May

    Under reviewPaper
  4. Understanding Communication Backends in Cross-Silo Federated Learning

    Amir Ziashahabi, Chaoyang He, Salman Avestimehr

    ICC 2026Paper
  5. Reject Only Critical Tokens: Pivot-Aware Speculative Decoding

    Amir Ziashahabi*, Yavuz Faruk Bakman*, Duygu Nur Yaldiz, Mostafa El-Khamy, Sai Praneeth Karimireddy, Salman Avestimehr

    NeurIPS 2025 Efficient Reasoning WorkshopPaperCode
  6. MobiZO: Enabling Efficient LLM Fine-Tuning at the Edge via Inference Engines

    Lei Gao*, Amir Ziashahabi*, Yue Niu, Salman Avestimehr, Murali Annavaram

    EMNLP 2025PaperCode
  7. GeoToken: Hierarchical Geolocalization of Images via Next Token Prediction

    Narges Ghasemi*, Amir Ziashahabi*, Salman Avestimehr, Cyrus Shahabi

    ICDM 2025PaperCode
  8. OSMGen: Highly Controllable Satellite Image Synthesis using OpenStreetMap Data

    Amir Ziashahabi*, Narges Ghasemi*, Sajjad Shahabi, John Krumm, Salman Avestimehr, Cyrus Shahabi

    NeurIPS 2025 UrbanAI WorkshopPaperCode
  9. Frequency Domain Diffusion Model with Scale-Dependent Noise Schedule

    Amir Ziashahabi, Baturalp Buyukates, Artan Sheshmani, Yi-Zhuang You, Salman Avestimehr

    ISIT 2024PaperCode
  10. Hawk: Accurate and Fast Privacy-Preserving Machine Learning Using Secure Lookup Table Computation

    Hamza Saleem, Amir Ziashahabi, Muhammad Naveed, Salman Avestimehr

    PETS 2024Paper
  11. Renormalization Group flow, Optimal Transport and Diffusion-based Generative Model

    Artan Sheshmani, Yi-Zhuang You, Baturalp Buyukates, Amir Ziashahabi, Salman Avestimehr

    Physical Review E 2024Paper
  12. PyTorch RPC: Distributed Deep Learning Built on Tensor-Optimized Remote Procedure Calls

    Shen Li, Pritam Damania, Luca Wehrstedt, et al., including Amir Ziashahabi

    MLSys 2023Paper

Experience

  • Machine Learning Researcher

    2021 – Present

    vITAL Lab, University of Southern California · Los Angeles, CA · Advisor: Prof. Salman Avestimehr

    • Conduct machine learning research with a focus on efficiency in system and algorithmic design.
    • Focus areas: large language models, federated learning, secure multiparty computation, distributed ML, diffusion models.
  • Machine Learning Intern

    2022

    FedML, Inc. · Los Angeles, CA

    • Designed and developed hierarchical federated learning systems and analyzed communication technologies performance.
  • Senior Developer & Distributed Systems Researcher

    2018 – 2021

    Rastak Media Sepehr Co. · Tehran, Iran

    • Led projects including a distributed ledger with Hyperledger Fabric and a ride‑hailing platform on Ethereum.
    • Engineered solutions using Kubernetes, Solidity, React, NestJS, Web3 APIs.
  • Distributed Systems Researcher

    2018 – 2020

    Decentralab Lab, Sharif University of Technology · Tehran, Iran · Advisor: Prof. Mohammad Ali Maddah‑Ali

    • Investigated scalability solutions for blockchain with a focus on layer‑2 solutions.
    • Thesis on a blockchain‑based social news aggregator; contributed to blockchain course curriculum.
  • Distributed Systems Researcher

    2016 – 2017

    Cloud Computing Center, Iran University of Science & Technology · Tehran, Iran · Advisor: Prof. Mehrdad Ashtiani

    • Researched IaaS cloud solutions; deployed and monitored OpenStack services.

Education

  • Ph.D. in Electrical Engineering

    2021 – 2026 (expected)

    University of Southern California · Los Angeles, CA · GPA 4/4

  • M.Sc. in Computer Engineering – Software

    2017 – 2020

    Sharif University of Technology · Tehran, Iran · GPA 19.38/20

  • B.Sc. in Computer Engineering – Software

    2013 – 2017

    Iran University of Science and Technology · Tehran, Iran · GPA 17.66/20

Teaching

  • Applied and Cloud Computing for Electrical EngineersTeaching Assistant, USC
  • Foundations of Blockchain TechnologyTeaching Assistant, Sharif University of Technology
  • Ethereum Programming WorkshopsInstructor, Various Locations
  • Iran Blockchain Summer SchoolCo‑organizer and Lecturer, Sharif University of Technology
  • Hands‑on Blockchain WorkshopCo‑organizer and Mentor, Sharif University of Technology
  • Database DesignTeaching Assistant, Sharif University of Technology
  • Compiler DesignTeaching Assistant, Iran University of Science and Technology

Reviewing

  • NeurIPS, ICML (Gold Reviewer Award 2026), AAAI
  • Efficient Reasoning @NeurIPS 2025, ENLSP @NeurIPS 2024, SCOPE @ICLR 2025

Other Projects

  • Software Fault Prediction. Designed and evaluated a concise fault‑prediction method on software‑evolution data. (R)
  • Toy Language + JVM Backend. Built a small language with a compiler front‑end and JVM bytecode backend. (ANTLR, Jasmin)
  • Hunterguh: 2D Multiplayer. Implemented a 2D multiplayer action game. (Unity, C#)
  • 2D Game Engine (from scratch). Lightweight engine with UI, input handling, and a scene/update loop. (Python, PyGame, PyQt)
  • Arcade Game. Fixed‑shooter arcade game implemented from scratch. (C++)
  • ChaapArt: Card‑Design Web App. Co‑founded and built the backend for a templated card‑design platform. (PHP, Laravel)

Skills

Programming
Python, JavaScript, Solidity, CSS, HTML, Go, C#, PHP, Java, C/C++
Tools and Technologies
PyTorch RPC, gRPC, MPI, Docker, Kubernetes, Redis, Nginx, Unity Game Engine, Truffle
Libraries/Frameworks
PyTorch, Transformers, Diffusers, vLLM, GPT-Fast, Express, NestJS, React, Django, Flask, Laravel