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. The mixinputs will not activate without the environment variable MIXINPUTS_BETA.

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.

Set MIXINPUTS_BETA=-100 to activate direct mixture, i.e., mixing with solely output distributions.

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).

For external evaluation packages, simply set export MIXINPUTS_BETA=<BETA_VALUE> and use vllm as the backbone, will activate MoI.

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.1.tar.gz (38.0 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.1-py3-none-any.whl (37.4 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: mixinputs-0.1.1.tar.gz
  • Upload date:
  • Size: 38.0 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.1.tar.gz
Algorithm Hash digest
SHA256 e7120ed5c9b01829dacba3cb9121e247297ac91f2df1c785db13df780a492384
MD5 05fdf8b919362e37719e12b6b5cc5fd4
BLAKE2b-256 582496b59eca86a18662d880dbd0c14a44396cc461254c8786a77fb9bae13e08

See more details on using hashes here.

File details

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

File metadata

  • Download URL: mixinputs-0.1.1-py3-none-any.whl
  • Upload date:
  • Size: 37.4 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.1-py3-none-any.whl
Algorithm Hash digest
SHA256 659836980b4f37d307460aa8be69358fca3ae9df65a1611c1f8af357708e5abf
MD5 da2c962b4939ac57dd3804144d28fda3
BLAKE2b-256 661e183555393d77f088ab50469d6f656fa70acff8a23e0d318eac62dc207d61

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