Simple and effective multi-token completion from masked language models

By Oren Kalinsky, Guy Kushilevitz, Alex Libov, Yoav Goldberg
2023
Download Copy BibTeX GitHub
Copy BibTeX
Pre-trained neural masked language models are often used for predicting a replacement token for a given sequence position, in a cloze-like task. However, this usage is restricted to predicting a single token, from a relatively small pre-trained vocabulary. Recent Sequence2Sequence pre-trained LMs like T5 do allow predicting multi-token completions, but are more expensive to train and run. We show that pre-trained masked language models can be adapted to produce multi-token completions, with only a modest addition to their parameter count. We propose two simple adaptation approaches, trading parameter counts for accuracy. The first method generates multitoken completions from a conditioned RNN. It has a very low parameter count and achieves competitive results. The second method is even simpler: it adds items corresponding to multi-token units to the output prediction matrix. While being higher in parameter count than the RNN method, it also surpasses current state-of-the-art multi-token completion models, including T5-3B, while being significantly more parameter efficient. We demonstrate that our approach is flexible to different vocabularies and domains and can effectively leverage existing pre-trained models available in different domains. Finally, a human evaluation further validates our results and shows that our solution regularly provides valid completions, as well as reasonable correctness for factual-sentence completions.
Research areas

Latest news

US, NY, New York
We are seeking a Human-Robot Interaction (HRI) Applied Scientist to develop cutting-edge interactions that make robots feel alive, personal, and fun. In this role, you will focus on verbal and non-verbal conversational systems, social dynamics, memory, and long-term relationship formation between robots, their environments, and the people they interact with. Your contributions will be essential in advancing robotics by enabling expressive, socially intelligent, and trustworthy interactions between robots and humans. Key job responsibilities - Develop interactive systems that leverage large language models, multimodal inputs and outputs, reinforcement learning from human feedback, or other advanced techniques to achieve fluid, engaging, and socially appropriate robot behavior - Design and implement intelligent conversational systems that handle turn-taking, grounding, interruption, and incorporates context drawn from a robot's physical environment and shared history with a user - Integrate perceptual sensor streams including gaze, facial expression, gesture, posture, and more to understand social context and produce coherent, lifelike interactions. - Develop memory and personalization systems that allow robots to form lasting relationships with individual users, learn their environments, and adapt their behavior over weeks and months - Stay updated on advancements in HRI, NLP, multimodal AI, and cognitive and social science to apply cutting-edge techniques to robot interaction challenges - Lead technical projects from conception through production deployment - Mentor junior scientists and engineers - Bridge research initiatives with practical engineering implementation
US, NY, New York
We are seeking a Human-Robot Interaction (HRI) Applied Scientist to develop cutting-edge interactions that make robots feel alive, personal, and fun. In this role, you will focus on verbal and non-verbal conversational systems, social dynamics, memory, and long-term relationship formation between robots, their environments, and the people they interact with. Your contributions will be essential in advancing robotics by enabling expressive, socially intelligent, and trustworthy interactions between robots and humans. Key job responsibilities - Develop interactive systems that leverage large language models, multimodal inputs and outputs, reinforcement learning from human feedback, or other advanced techniques to achieve fluid, engaging, and socially appropriate robot behavior - Design and implement intelligent conversational systems that handle turn-taking, grounding, interruption, and incorporates context drawn from a robot's physical environment and shared history with a user - Integrate perceptual sensor streams including gaze, facial expression, gesture, posture, and more to understand social context and produce coherent, lifelike interactions. - Develop memory and personalization systems that allow robots to form lasting relationships with individual users, learn their environments, and adapt their behavior over weeks and months - Stay updated on advancements in HRI, NLP, multimodal AI, and cognitive and social science to apply cutting-edge techniques to robot interaction challenges - Lead technical projects from conception through production deployment - Mentor junior scientists and engineers - Bridge research initiatives with practical engineering implementation
US, CA, San Francisco
Amazon is on a mission to redefine the future of automation — and we're looking for exceptional talent to help lead the way. We are building the next generation of advanced robotic systems that seamlessly blend cutting-edge AI, sophisticated control systems, and novel mechanical design to create adaptable, intelligent automation solutions capable of operating safely alongside humans in dynamic, real-world environments. At Amazon, we leverage the power of machine learning, artificial intelligence, and advanced robotics to solve some of the most complex operational challenges at a scale unlike anywhere else in the world. Our fleet of robots spans hundreds of facilities globally, working in sophisticated coordination to deliver on our promise of customer excellence — and we're just getting started. As a Scientist in Robot Navigation, you will be at the forefront of this transformation — architecting and delivering navigation systems that are intelligent, safe, and scalable. You will bring deep expertise in learning-based planning and control, a strong understanding of foundation models and their application to embodied agents, and as well as have in-depth understanding of control-theoretic approaches such as model predictive control (MPC)-based trajectory planning. You will develop navigation solutions that seamlessly blend data-driven intelligence with principled control-theoretic guarantees. Our vision is bold: to build navigation systems that allow robots to move fluidly and safely through dynamic environments — understanding context, anticipating change, and adapting in real time. You will lead research that bridges the gap between cutting-edge academic advances and production grade deployment, collaborating with world-class teams pushing the boundaries of robotic autonomy, manipulation, and human-robot interaction. Join us in building the next generation of intelligent navigation systems that will define the future of autonomous robotics at scale. Key job responsibilities - Design, develop, and deploy perception algorithms for robotics systems, including object detection, segmentation, tracking, depth estimation, and scene understanding - Lead research initiatives in computer vision, sensor fusion and 3D perception - Collaborate with cross-functional teams including robotics engineers, software engineers, and product managers to define and deliver perception capabilities - Drive end-to-end ownership of ML models — from data collection and labeling strategy to training, evaluation, and deployment - Mentor junior scientists and engineers; contribute to a culture of technical excellence - Define and track key metrics to measure perception system performance in real-world environments - Publish research findings in top-tier venues (CVPR, ICCV, ECCV, ICRA, NeurIPS, etc.) and contribute to patents A day in the life - Train ML models for deployment in simulation and real-world robots, identify and document their limitations post-deployment - Drive technical discussions within your team and with key stakeholders to develop innovative solutions to address identified limitations - Actively contribute to brainstorming sessions on adjacent topics, bringing fresh perspectives that help peers grow and succeed — and in doing so, build lasting trust across the team - Mentor team members while maintaining significant hands-on contribution to technical solutions About the team Our team is a group is a diverse group of scientists and engineers passionate about building intelligent machines. We value curiosity, rigor, and a bias for action. We believe in learning from failure and iterating quickly toward solutions that matter.
IN, KA, Bengaluru
Every product a customer returns is a moment where Amazon either recovers value or writes it off — and India's ReCommerce business is on a multi-million-dollar mission to recover more of it, more intelligently, at scale. Machine learning is the core lever: predicting whether a returned unit is sellable without a human touching it, detecting damage and fraud inside sealed packaging from images, routing each unit to its highest-value disposition, and pricing recovered inventory dynamically. India's returns network is large, fast-growing, and structurally different from other geographies — a rich, high-impact environment for an Applied Scientist to build models that move real financial and customer-experience metrics. We are hiring an Applied Scientist to build and adapt the ML that powers India ReCommerce. You will work at the intersection of two mandates: building India-first models for problems unique to our market, and adapting proven Worldwide models to India's data, catalog, and operational reality — recalibrating them where distribution, language, and process differ. You will own problems end-to-end, from framing and data through modeling, evaluation, and production deployment, partnering closely with engineering, product, and operations. Key job responsibilities Build ML models for automated returns grading — predicting the salability of returned units from structured and unstructured signals so units can be evaluated with zero or minimal human touch, improving speed, accuracy, and recovery value. Develop computer-vision models for defect detection, condition assessment, and anomaly/fraud identification (including inside sealed packaging), and for establishing chain-of-custody and damage attribution across the returns journey. Build disposition-prediction and routing models that direct each unit to its highest-value recovery path (resale, repair, liquidation, donation, recycle) as early as possible in the network. Develop pricing and recovery-optimization models for liquidation and resale, moving from flat rates toward dynamic, grade- and condition-aware pricing. Adapt Worldwide ML models to India — retraining, recalibrating, and re-evaluating for India's return distribution, catalog, languages, and operational constraints, and closing the gaps that prevent a direct lift-and-shift. Own the full model lifecycle — problem framing, data pipelines, feature engineering, training, offline/online evaluation, monitoring, and retraining — with rigorous attention to calibration, drift, and business-metric impact. Partner cross-functionally with engineering (to productionize), product (to frame problems and measure impact), and operations (to ground models in how the network actually runs), and use modern GenAI/LLM tooling to accelerate research and delivery. A day in the life You start by reviewing the performance of a grading model in production — checking calibration and drift against last week's returns, and confirming the recovery-value lift is holding. Mid-morning, you dig into a computer-vision problem: improving detection of a damage type that's driving write-offs, using images captured across the returns journey. In the afternoon you work with a Worldwide science team to bring one of their models to India — scoping what retraining and recalibration India's data requires — then pair with an engineer to move your latest model toward production behind a clean evaluation gate. You close by framing a new problem with a product partner: quantifying the opportunity, defining the label and success metric, and sketching the modeling approach. About the team India ReCommerce owns the systems and science that turn returned and unsellable inventory into recovered value and a better customer experience. You will join a team building an increasingly automated, ML-driven returns network — leveraging Worldwide platforms where they fit and building India-first capabilities where they don't. It is a high-ownership environment with a direct line from your models to measurable business and customer outcomes.
IN, MH, Mumbai
Amazon Science gives you insight into the company’s approach to customer-obsessed scientific innovation. Amazon fundamentally believes that scientific innovation is essential to being the most customer-centric company in the world. It’s the company’s ability to have an impact at scale that allows us to attract some of the brightest minds in artificial intelligence and related fields. Our scientists continue to publish, teach, and engage with the academic community, in addition to utilizing our working backwards method to enrich the way we live and work. Please visit https://www.amazon.science for more information. About Amazon Prime Video “Many of the problems we face have no textbook solution, and so we-happily-invent new ones.” – Jeff Bezos
 The Amazon Prime Video team is shaping the future of digital video entertainment. We are seeking a Data Scientist to uncover key insights on how consumers watch videos on Amazon. The ideal candidate will be an expert in the areas of data science, machine learning and statistics, having hands-on experience with multiple improvement initiatives as well as balancing technical and business judgment to make the right decisions about technology, models and methodologies. As consumers increasingly consume digital video, we need to make agile decisions based on what content appeals to our customers. As a Data Scientist at Amazon Prime Video APAC and ANZ analytics team, you will have the opportunity to work on one of the world's largest consumer data sets, influence the long term evolution of our analytics capability and support the expansion of Amazon's digital video business. The Data Scientist will work closely with other research scientists, machine-learning experts, and economists to design and run experiments, research new algorithms, and find new ways to improve optimization across all our associate facing tools. 
 A successful candidate will be able to understand and manage key operational and technical concepts. They will have excellent project and communication skills, and motivation to achieve results in a fast-paced environment. Candidates should demonstrate a passion for working on behalf of customers, have a record of accomplishment of timely delivery of large-scale projects, and have the ability to influence multiple global teams. Autonomy, judgment, influence, and leadership skills are essential. This person will be responsible for ensuring we meet our key deliverables, on time with high quality, and communicating status to internal and external stakeholders. Key Responsibilities - Support the Content team on business reporting, ad hoc analysis, statistical inference and predictive modelling for all Prime Video APAC and ANZ. - Mine and analyze data pertaining to customers viewing experiences to identify critical business insight and make recommendations to optimize content selection. - Proactively develop new ML models using streaming, video, audio and textual data to understand and predict customer streaming behaviour - Translate analytic insights into concrete, actionable recommendations for business or product improvement. Develop and present these as papers to senior stakeholders. - Liaise with your peers in other prime video territories to develop solutions that greatly benefit our global customers - This role will be based in Mumbai, India
US, WA, Seattle
Join us at the forefront of Amazon's sustainability initiatives to work on environmental and social advancements that support Amazon's long-term worldwide sustainability strategy. At Amazon, we're working to be the most customer-centric company on earth. To get there, we need exceptionally talented, bright, and driven people. We are looking for a Senior Research Scientist to join our growing Sustainability team to drive the science behind value chain decarbonization. This role will establish Amazon's scientific methodologies for sector- and cross-sectoral decarbonization mechanisms and establish benchmarks for automated validation and risk assessment. As a Senior Research Scientist, you will be responsible for independently leading assessments of environmental issues across the full spectrum of Amazon businesses and evaluating sustainability impacts across the value chain. You will independently develop quality frameworks and methodologies that enable Amazon to scale procurement of high-quality environmental interventions while maintaining scientific rigor and environmental integrity. Key job responsibilities - Develop quality assessment frameworks for complex environmental interventions, baseline-setting approaches, and measurement methodologies - Build quantitative benchmark and statistical models that enable scalable evaluation across heterogeneous data sources - Create attribution methodologies for supply chain interventions across Amazon's diverse footprint - Develop social and environmental safeguard criteria that integrate community impact assessments - Collaborate with cross-functional teams including procurement, sustainability operations, and business units to translate scientific methodologies into operational requirements - Work under the direction of senior business leaders while acting as lead Subject Matter Expert for value chain decarbonization science, including designing and leading research, data collection, modeling, documentation, interpretation, and validation About the team Diverse Experiences: Worldwide Sustainability values diverse experiences. Even if you do not meet all of the qualifications and skills listed in the job description, we encourage candidates to apply. If your career is just starting, hasn’t followed a traditional path, or includes alternative experiences, don’t let it stop you from applying. Inclusive Team Culture: It’s in our nature to learn and be curious. Our employee-led affinity groups foster a culture of inclusion that empower us to be proud of our differences. Ongoing events and learning experiences, including our Conversations on Race and Ethnicity (CORE) and AmazeCon (inclusive diversity) conferences, inspire us to never stop embracing our uniqueness. Mentorship & Career Growth: We’re continuously raising our performance bar as we strive to become Earth’s Best Employer. That’s why you’ll find endless knowledge-sharing, mentorship and other career-advancing resources here to help you develop into a better-rounded professional.
US, WA, Redmond
Amazon Leo is Amazon’s low Earth orbit satellite broadband network. Its mission is to deliver fast, reliable internet to customers and communities around the world, and we’ve designed the system with the capacity, flexibility, and performance to serve a wide range of customers, from individual households to schools, hospitals, businesses, government agencies, and other organizations operating in locations without reliable connectivity. Export Control Requirement: Due to applicable export control laws and regulations, candidates must be a U.S. citizen or national, U.S. permanent resident (i.e., current Green Card holder), or lawfully admitted into the U.S. as a refugee or granted asylum. We are looking for an experienced Data Scientist to help architect state-of-the-art test infrastructure and lead the development of data models and analysis tools to represent the ground truth about satellite test results in order to facilitate important business decisions. Our team is responsible for core infrastructure and tools that will serve as the backbone of automated satellite testing operations to enable rapid scaling of manufacturing processes. Key job responsibilities * Work with engineering, software and manufacturing teams to understand drivers, impacts, and key influences on satellite performance * Lead the design, build and implementation of production models and make decisions in real time for satellite test results * Drive actions at scale to optimize test methodology and drive increases to satellite reliability * Analysis and modeling of satellite telemetry from test results in lab and on-orbit * Develop models and data pipelines for satellite telemetry * Create and manage datasets for continued pre-training and supervised fine-tuning of LLMs * Develop scalable visualizations for analysis of satellite performance A day in the life As Amazon Leo Data Scientist you will own the architecture definition and development of data analysis tools to to aid engineering and production teams in deciding flight-worthiness of each Amazon Leo satellite and historical traceability tools to enable simplified discovery and interpretation of past test data. You will work with multiple engineering, software and manufacturing teams across ground and space systems, to specify requirements, define data collection, interpretation strategies, data pipelines and implement data analysis and reporting tools for Integrated Vehicle tests. Your focus will be in optimizing the analysis of test results to enable Amazon Leo production plans. About the team The Automated Vehicle Testing Team is a mix of scientists and software engineers responsible for data infrastructure, tools, and research that serve as the backbone of automated satellite testing operations to enable rapid scaling of manufacturing processes.
CN, 31, Shanghai
Worldwide Global Selling has been helping individuals and businesses increase sales and reach new customers around the globe. Today, more than 50% of Amazon's total unit sales come from third-party selection. The Global Selling team in China is responsible for recruiting local businesses to sell on Amazon's 19+ overseas marketplaces and supporting local Sellers' success and growth on Amazon. Our vision is to be the first choice for all types of Chinese business to go globally. The Worldwide Global Selling Analytics, Intelligence, and Technology (WWGS-AIT) team serves as the research, automation, and insight arm of the International Seller Service data hub, enabling rapid delivery of growth insights through strategic investments in regional data foundations, self-service business intelligence solutions, and artificial intelligence tools. The WWGS-AIT team is positioned to establish AI-ready foundational capabilities across the WWGS organization while maintaining excellence in business insight generation, and self-service BI/AI application development. WWGS-AIT is looking for a Data Scientist to design and build seller-facing AI agents that turn our AI-ready data foundation into intelligent, conversational experiences for Amazon's global sellers. You will own the intelligence layer of these agents end-to-end, from modeling and retrieval to evaluation and launch, working alongside applied scientists, data engineers, and the Seller Assistant platform team to put trustworthy AI directly into sellers' hands. Key job responsibilities - Design, build, and iterate seller-facing AI agents (LLM-powered) that help Chinese sellers grow globally, reasoning over WWGS-AIT's AI-ready data foundation and knowledge base. - Develop the intelligence layer of agents: retrieval-augmented generation (RAG) over our knowledge management system, tool-use / function-calling orchestration, prompt engineering, and model fine-tuning or adaptation where needed. - Ground agent responses in standardized metrics and unified seller profiles to guarantee consistency and accuracy across agents; design and enforce guardrails that prevent hallucination and protect sensitive, compliance-restricted data. - Build rigorous evaluation frameworks (golden datasets, offline evaluation, and online experimentation) to measure and continuously improve agent quality, safety, and seller impact. - Develop seller-intelligence models (segmentation, entity resolution / One-ID, ranking and recommendation) that power personalized agent experiences. - Partner with WWGS Tech and the Seller Assistant platform team to productionize agents and tools (e.g., via MCP), defining the model and intelligence contract while engineering operates the runtime. - Collaborate with business, product, and cross-functional partners to translate seller pain points into agent capabilities and measurable business outcomes. - Stay current with advances in GenAI and agentic systems, and bring applied research into production.
US, WA, Seattle
We are looking for an seasoned Applied Scientist to design, build, and deploy the ML/AI models that power our decision intelligence platform. You will work at the intersection of causal inference, time-series forecasting, anomaly detection, and LLM-based reasoning — all applied to real operational problems with measurable business impact. Key Job Responsibilities Decision Intelligence Models - **Causal inference & root cause analysis:** Build models that decompose fleet-wide metric movements into root causes, distinguishing correlation from causation across operational dimensions (site, service, failure mode, time) - **Dose-response modeling:** Develop models that learn the quantitative relationship between intervention intensity and outcome magnitude - **Forecasting & projection:** Build time-series models that project metric trajectories under different intervention scenarios, enabling "if we do X, expect Y by date Z" recommendations - **Anomaly detection & trend identification:** Develop multi-variate anomaly detection that distinguishes signal from noise in noisy operational data, and identifies emerging patterns before they become crises - **Confidence calibration:** Build and maintain calibrated confidence scores for recommendations, ensuring the system knows what it knows and what it doesn't - **Outcome attribution:** Design experiments and causal methods to measure the true impact of interventions LLM Integration & Reasoning - **Structured reasoning:** Design LLM prompting architectures that reliably transform operational data into executive-quality narrative summaries, decision framings, and recommendation rationales - **LLM evaluation:** Build evaluation frameworks that measure LLM output quality (accuracy, actionability, calibration) and detect degradation over time - **RAG systems:** Design retrieval-augmented generation systems that ground LLM outputs in operational data, historical playbooks, and institutional knowledge - **Progressive autonomy:** Design the trust-calibration system where AI gradually earns expanded authority based on demonstrated accuracy over time Research & Production - **End-to-end ownership:** Take models from research through production deployment — you ship, you monitor, you iterate - **Experimentation:** Design A/B tests and quasi-experiments to validate model improvements and measure business impact - **Stakeholder communication:** Translate complex scientific results into actionable insights for non-technical senior leaders
US, CA, Sunnyvale
Prime Video is a first-stop entertainment destination offering customers a vast collection of premium programming in one app available across thousands of devices. Prime members can customize their viewing experience and find their favorite movies, series, documentaries, and live sports – including Amazon MGM Studios-produced series and movies; licensed fan favorites; and programming from Prime Video subscriptions such as Apple TV+, HBO Max, Peacock, Crunchyroll and MGM+. All customers, regardless of whether they have a Prime membership or not, can rent or buy titles via the Prime Video Store, and can enjoy even more content for free with ads. Are you interested in shaping the future of entertainment? Prime Video's technology teams are creating best-in-class digital video experience. As a Prime Video team member, you’ll have end-to-end ownership of the product, user experience, design, and technology required to deliver state-of-the-art experiences for our customers. You’ll get to work on projects that are fast-paced, challenging, and varied. You’ll also be able to experiment with new possibilities, take risks, and collaborate with remarkable people. We’ll look for you to bring your diverse perspectives, ideas, and skill-sets to make Prime Video even better for our customers. With global opportunities for talented technologists, you can decide where a career Prime Video Tech takes you! Key job responsibilities As an Applied Scientist at Prime Video, you will have end-to-end ownership of the product, related research and experimentation, applying advanced machine learning techniques in computer vision (CV), Generative AI, multimedia understanding and so on. You’ll work on diverse projects that enhance Prime Video’s content localization, image/video understanding, and content personalization, driving impactful innovations for our global audience. Other responsibilities include: - Research and develop generative models for controllable synthesis across images, video, vector graphics, and multimedia - Innovate in advanced diffusion and flow-based methods (e.g., inverse flow matching, parameter efficient training, guided sampling, test-time adaptation) to improve efficiency, controllability, and scalability. - Advance visual grounding, depth and 3D estimation, segmentation, and matting for integration into pre-visualization, compositing, VFX, and post-production pipelines. - Design multimodal GenAI workflows including visual-language model tooling, structured prompt orchestration, agentic pipelines. A day in the life Prime Video is pioneering the use of Generative AI to empower the next generation of creatives. Our mission is to make world-class media creation accessible, scalable, and efficient. We are seeking an Applied Scientist to advance the state of the art in Generative AI and to deliver these innovations as production-ready systems at Amazon scale. Your work will give creators unprecedented freedom and control while driving new efficiencies across Prime Video’s global content and marketing pipelines. This is a newly formed team within Prime Video Science!