Publications

2023

1. AAAI

   PiCor: Multi-Task Deep Reinforcement Learning with Policy Correction

   Fengshuo Bai, Hongming Zhang, Tianyang Tao, Zhiheng Wu, Yanna Wang, and Bo Xu

   In AAAI Conference on Artificial Intelligence, 2023

   Abs PDF

   Multi-task deep reinforcement learning (DRL) ambitiously aims to train a general agent that masters multiple tasks si- multaneously. However, varying learning speeds of different tasks compounding with negative gradient interference makes policy learning inefficient. In this work, we propose PiCor, an efficient multi-task DRL framework that splits learning into policy optimization and policy correction phases. The policy optimization phase improves the policy by any DRL algo- thrim on the sampled single task without considering other tasks. The policy correction phase first constructs a perfor- mance constraint set with adaptive weight adjusting. Then the intermediate policy learned by the first phase is constrained to the set, which controls the negative interference and balances the learning speeds across tasks. Empirically, we demonstrate that PiCor outperforms previous methods and significantly improves sample efficiency on simulated robotic manipula- tion and continuous control tasks. We additionally show that adaptive weight adjusting can further improve data efficiency and performance.

2024

1. preprint

   Provably Robust Multi-bit Watermarking for AI-generated Text via Error Correction Code

   Wenjie Qu, Dong Yin, Zixin He, Wei Zou, Tianyang Tao, Jinyuan Jia, and Jiaheng Zhang

   In arXiv, 2024

   Abs PDF

   Large Language Models (LLMs) have been widely deployed for their remarkable capability to generate texts resembling human language. However, they could be misused by criminals to create deceptive content, such as fake news and phishing emails, which raises ethical concerns. Watermarking is a key technique to mitigate the misuse of LLMs, which embeds a watermark (e.g., a bit string) into a text generated by a LLM. Consequently, this enables the detection of texts generated by a LLM as well as the tracing of generated texts to a specific user. The major limitation of existing watermark techniques is that they cannot accurately or efficiently extract the watermark from a text, especially when the watermark is a long bit string. This key limitation impedes their deployment for real-world applications, e.g., tracing generated texts to a specific user. This work introduces a novel watermarking method for LLM-generated text grounded in \textbferror-correction codes to address this challenge. We provide strong theoretical analysis, demonstrating that under bounded adversarial word/token edits (insertion, deletion, and substitution), our method can correctly extract watermarks, offering a provable robustness guarantee. This breakthrough is also evidenced by our extensive experimental results. The experiments show that our method substantially outperforms existing baselines in both accuracy and robustness on benchmark datasets. For instance, when embedding a bit string of length 12 into a 200-token generated text, our approach attains an impressive match rate of 98.4%, surpassing the performance of Yoo et al. (state-of-the-art baseline) at 85.6%. When subjected to a copy-paste attack involving the injection of 50 tokens to generated texts with 200 words, our method maintains a substantial match rate of 90.8%, while the match rate of Yoo et al. diminishes to below 65%.

2. preprint

   An Efficient and Extensible Zero-knowledge Proof Framework for Neural Networks

   Tao Lu, Haoyu Wang, Wenjie Qu, Zonghui Wang, Jinye He, Tianyang Tao, Wenzhi Chen, and Jiaheng Zhang

   In Cryptology ePrint, 2024

   Abs PDF

   In recent years, cloud vendors have started to supply paid services for data analysis by providing interfaces of their well-trained neural network models. However, customers lack tools to verify whether outcomes supplied by cloud vendors are correct inferences from particular models, in the face of lazy or malicious vendors. The cryp- tographic primitive called zero-knowledge proof (ZKP) addresses this problem. It enables the outcomes to be verifiable without leak- ing information about the models. Unfortunately, existing ZKP schemes for neural networks have high computational overheads, especially for the non-linear layers in neural networks. In this paper, we propose an efficient and extensible ZKP frame- work for neural networks. Our work improves the performance of the proofs for non-linear layers. Compared to previous works relying on the technology of bit decomposition, we convert com- plex non-linear relations into range and exponent relations, which significantly reduces the number of constraints required to prove non-linear layers. Moreover, we adopt a modular design to make our framework compatible with more neural networks. Specifically, we propose two enhanced range and lookup proofs as basic blocks. They are efficient in proving the satisfaction of range and exponent relations. Then, we constrain the correct calculation of primitive non-linear operations using a small number of range and exponent relations. Finally, we build our ZKP framework from the primitive operations to the entire neural networks, offering the flexibility for expansion to various neural networks. We implement our ZKPs for convolutional and transformer neu- ral networks. The evaluation results show that our work achieves over 168.6× (up to 477.2×) speedup for separated non-linear layers and 41.4× speedup for the entire ResNet-101 convolutional neural network, when compared with the state-of-the-art work, Mystique. In addition, our work can prove GPT-2, a transformer neural net- work with 117 million parameters, in 287.1 seconds, achieving 35.7× speedup over ZKML, which is a state-of-the-art work supporting transformer neural networks.