Skip to main content

Mixture of Inputs Patch for vLLM

Project description

🎨 Mixture of Inputs (MoI) 🎨

Text Generation Beyond Discrete Token Sampling (arXiv)

MoI Illustration


What is Mixture of Inputs?

Mixture of Inputs (MoI) is a simple, training-free technique to improve autoregressive text generation.

In standard LLM inference, the model samples a token and discards the full distribution. MoI keeps both—it mixes the sampled token with the original distribution to retain more information.

MoI uses Bayesian estimation to compute mixing weights:

  • Prior: the output distribution

  • Observation: the sampled token

  • Posterior: a smooth input replacing the one-hot vector

This lets the model use a richer representation state $\boldsymbol{h}_t$ as inputs throughout generation, improving coherence, reasoning, and code synthesis.


Installation

Requires vllm >= 0.8.5.

pip install vllm
pip install mixinputs
mixinputs setup

Quick Start

To activate MoI, set the MIXINPUTS_BETA environment variable:

export MIXINPUTS_BETA=1.0

Then run your usual vLLM-based generation script. That's it!

CLI Utilities

The patch is installed via injection to usercustomize.py by adding import mixinputs at the head.

We provide command-line tools to patch or unpatch your environment:

# Enable MoI
mixinputs setup

# Disable MoI
mixinputs cleanup

You can also enable it via adding import mixinputs in your script before you import vllm.

Configuration Options

Variable Description Default
MIXINPUTS_BETA Controls the strength of the mixture input signal 1.0

Recommended range: 0.5 to 2.0.

Tune based on task/model—lower values emphasize the distribution, higher values keep more of the sample.

Make sure enforce_eager=True in your LLM initialization.

Examples

We have included 2 example usages with AIME and Count Down 4, after executing the bash you should see 59-60 Acc for CountDown4 (Nemotron-Super-49B) and ~80 Acc for AIME (QwQ-32B).

Additionally, MoI also work with server mode, we provide an example script.

Evaluations

Model Method Input Info. AIME (%) CountDown4 (%) GPQA-D (%) LiveCodeBench (pass@1) Avg (%)
QwQ-32B Standard Output Token 77.78 79.25 58.08 76.32 72.86
Direct Mixture Output Dist. 72.00 66.88 51.52 53.42 60.96
MoI Token + Dist. 80.00 80.01 60.10 76.51 74.15
Gain vs. Standard +2.22 +0.76 +2.02 +0.19 +1.29
Nemotron-Super-49B Standard Output Token 54.89 56.93 60.60 39.92 53.09
Direct Mixture Output Dist. 60.00 51.72 60.10 16.04 46.97
MoI Token + Dist. 57.11 59.53 64.65 40.50 55.45
Gain vs. Standard +2.22 +2.60 +4.05 +0.58 +2.36
Gemma-3-27B Standard Output Token 25.56 56.51 46.97 31.31 40.09
Direct Mixture Output Dist. 26.44 55.47 51.52 31.99 41.36
MoI Token + Dist. 26.89 59.38 47.47 32.87 41.65
Gain vs. Standard +1.33 +2.87 +0.50 +1.56 +1.56
DAPO-Qwen-32B Standard Output Token 64.67 72.03 42.42 54.01 58.28
Direct Mixture Output Dist. 62.67 67.19 37.88 23.87 47.90
MoI Token + Dist. 64.44 78.75 42.93 55.18 60.33
Gain vs. Standard –0.23 +6.72 +0.51 +1.17 +2.05

Questions?

If you have any questions related to the code or the paper, feel free to reach out to us at y5zhuang@ucsd.edu.

Citation

If you find our paper and code useful, please cite us:

@article{zhuang2025textgen,
  title={Text Generation Beyond Discrete Token Sampling},
  author={Zhuang, Yufan and Liu, Liyuan and Singh, Chandan and Shang, Jingbo and Gao, Jianfeng},
  journal={arXiv preprint arXiv:2505.14827},
  year={2025}
}

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

mixinputs-0.1.0.tar.gz (37.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

mixinputs-0.1.0-py3-none-any.whl (37.3 kB view details)

Uploaded Python 3

File details

Details for the file mixinputs-0.1.0.tar.gz.

File metadata

  • Download URL: mixinputs-0.1.0.tar.gz
  • Upload date:
  • Size: 37.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for mixinputs-0.1.0.tar.gz
Algorithm Hash digest
SHA256 3587220a22ea6c5d7d3b58a339e15b2c87e6350f2a602b466117c9a06d9a362e
MD5 03e800377b5bc7a0f280eccf18c2ccd1
BLAKE2b-256 06397a8496e2d8e993b9b145ca17ae3a6e0225ad9f64c73897b6349afea998e1

See more details on using hashes here.

File details

Details for the file mixinputs-0.1.0-py3-none-any.whl.

File metadata

  • Download URL: mixinputs-0.1.0-py3-none-any.whl
  • Upload date:
  • Size: 37.3 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for mixinputs-0.1.0-py3-none-any.whl
Algorithm Hash digest
SHA256 c2afbac9c06b2ec07ed2da21416729b6b0d116829733ac8d8502d7b92a5a847b
MD5 4374bbeeeed1523d9833990ce932f7e3
BLAKE2b-256 8f9b4e755eb3bc116cda56e1b100a7d2c96fadfd6708317d951fc86c403861e2

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page