File size: 3,976 Bytes
106b238
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
---
license: apache-2.0
---
# Make Some Noise (MSN) Framework
Implementation of EMNLP 2024 paper [Make Some Noise: Unlocking Language Model Parallel Inference
Capability through Noisy Training](https://arxiv.org/pdf/2406.17404).
[[Github]](https://github.com/wyxstriker/MakeSomeNoiseInference)

## Requirements
- Environment: We adopt the same environment as used in [Spec-Bench](https://github.com/hemingkx/Spec-Bench) to facilitate a fair and consistent evaluation.
- Prepared Models: For convenience of testing, we release the weights of both the general-purpose model [[Llama3-8B-MSN](https://huggingface.co/DecoderImmortal/Llama3-8B-MSN)] and the code-specific model [[DeepSeek-Coder-7B-MSN](https://huggingface.co/DecoderImmortal/DeepSeek-Coder-7B-MSN)] trained on MSN as discussed in the paper.

## A minimal implementation of MSN

The MSN framework can be easily integrated into the data preprocessing stage of any training script. The entire noise addition process is as follows:

```python
# L denotes the noise length hyperparameter, which is typically set to 5.
dataset = [
    {"source_ids": "Query prompt.",
    "input_ids": "Concatenation of the query and response.",
    "output_ids": "Copy of input_ids as label for LM task."}
]
for source_ids, input_ids in dataset:
    start_idx = random.randrange(len(source_ids), len(input_ids)-L)
    for mask_i in range(start_idx, start_idx+L):
        # Noise is added only to the input portion corresponding to the response.
        input_ids[mask_i] = random.choice(input_ids[:mask_i])

```

## TR-Jacobi

<div align="center">
<img src="./pic/tr-jacobi.png" width="50%"/>
</div>

We demonstrate how to use TR-Jacobi to accelerate the MSN-trained model in ```src/inference_msn.py```.

```python
# jacobi decoding
spec_res_ids, new_tokens, forward_steps, accpet_list = noise_forward(input_ids.cuda(), model, tokenizer, args.max_new_tokens)

print("msn output")
print(tokenizer.decode(spec_res_ids[0]))
print("#MTA")
print(new_tokens/forward_steps)
print("Accepted Length List")
print(accpet_list)

# msn output
# <|begin_of_text|><|start_header_id|>system<|end_header_id|>
# Give me some advices about how to write an academic paper?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
# 1. Start by researching your topic and gathering relevant information. Make sure to take notes and organize your research in a way that makes sense.
# ...
# 8. Submit your paper. Make sure to follow any submission guidelines and make sure to submit your paper on time.<|eot_id|><|eot_id|>.

# #MTA
# 2.2

# Accepted Length List
# [1, 2, 1, 1, 3, 1, 2, 2, 3, 1, 2, 2, 2, 2, 2, 1, 3, 1, 3, 1, 2, 1, 3, 2, 2, 2, 1, 2, 1, 2, 3, 2, 3, 3, 2, 2, 2, 2, 2, 2, 2, 2, 1, 5, 1, 3, 1, 5, 2, 1, 3, 2, 2, 2, 3, 2, 5, 1, 3, 2, 3, 2, 3, 2, 1, 4, 3, 1, 2, 2, 3, 6, 1, 2, 2, 2, 3, 2, 2, 3, 3, 2, 3, 2, 2, 2, 1, 2, 2, 2, 3, 3, 3, 1, 4, 2, 1, 2, 2, 2]
```

Run ```sh run_case.sh``` to obtain the execution process of a test sample.
The interface design of the entire ```noise_forward``` is kept consistent with Spec-Bench.







## Citation
If you find this work is useful for your research, please cite our paper:

```
@inproceedings{wang-etal-2024-make,
    title = "Make Some Noise: Unlocking Language Model Parallel Inference Capability through Noisy Training",
    author = "Wang, Yixuan  and
      Luo, Xianzhen  and
      Wei, Fuxuan  and
      Liu, Yijun  and
      Zhu, Qingfu  and
      Zhang, Xuanyu  and
      Yang, Qing  and
      Xu, Dongliang  and
      Che, Wanxiang",
    editor = "Al-Onaizan, Yaser  and
      Bansal, Mohit  and
      Chen, Yun-Nung",
    booktitle = "Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing",
    month = nov,
    year = "2024",
    address = "Miami, Florida, USA",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2024.emnlp-main.718/",
    doi = "10.18653/v1/2024.emnlp-main.718",
    pages = "12914--12926",
}
```