Exploiting Code Symmetries for Learning Program Semantics

Kexin Pei, Weichen Li, Qirui Jin, Shuyang Liu, Scott Geng, Lorenzo Cavallaro, Junfeng Yang, Suman Jana
Proceedings of the 41st International Conference on Machine Learning, PMLR 235:40092-40113, 2024.

Abstract

This paper tackles the challenge of teaching code semantics to Large Language Models (LLMs) for program analysis by incorporating code symmetries into the model architecture. We introduce a group-theoretic framework that defines code symmetries as semantics-preserving transformations, where forming a code symmetry group enables precise and efficient reasoning of code semantics. Our solution, SymC, develops a novel variant of self-attention that is provably equivariant to code symmetries from the permutation group defined over the program dependence graph. SymC obtains superior performance on five program analysis tasks, outperforming state-of-the-art code models, including GPT-4, without any pre-training. Our results suggest that code LLMs that encode the code structural prior via the code symmetry group generalize better and faster.

Cite this Paper

BibTeX
@InProceedings{pmlr-v235-pei24b, title = {Exploiting Code Symmetries for Learning Program Semantics}, author = {Pei, Kexin and Li, Weichen and Jin, Qirui and Liu, Shuyang and Geng, Scott and Cavallaro, Lorenzo and Yang, Junfeng and Jana, Suman}, booktitle = {Proceedings of the 41st International Conference on Machine Learning}, pages = {40092--40113}, year = {2024}, editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, volume = {235}, series = {Proceedings of Machine Learning Research}, month = {21--27 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/pei24b/pei24b.pdf}, url = {https://proceedings.mlr.press/v235/pei24b.html}, abstract = {This paper tackles the challenge of teaching code semantics to Large Language Models (LLMs) for program analysis by incorporating code symmetries into the model architecture. We introduce a group-theoretic framework that defines code symmetries as semantics-preserving transformations, where forming a code symmetry group enables precise and efficient reasoning of code semantics. Our solution, SymC, develops a novel variant of self-attention that is provably equivariant to code symmetries from the permutation group defined over the program dependence graph. SymC obtains superior performance on five program analysis tasks, outperforming state-of-the-art code models, including GPT-4, without any pre-training. Our results suggest that code LLMs that encode the code structural prior via the code symmetry group generalize better and faster.} }
Endnote
%0 Conference Paper %T Exploiting Code Symmetries for Learning Program Semantics %A Kexin Pei %A Weichen Li %A Qirui Jin %A Shuyang Liu %A Scott Geng %A Lorenzo Cavallaro %A Junfeng Yang %A Suman Jana %B Proceedings of the 41st International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Ruslan Salakhutdinov %E Zico Kolter %E Katherine Heller %E Adrian Weller %E Nuria Oliver %E Jonathan Scarlett %E Felix Berkenkamp %F pmlr-v235-pei24b %I PMLR %P 40092--40113 %U https://proceedings.mlr.press/v235/pei24b.html %V 235 %X This paper tackles the challenge of teaching code semantics to Large Language Models (LLMs) for program analysis by incorporating code symmetries into the model architecture. We introduce a group-theoretic framework that defines code symmetries as semantics-preserving transformations, where forming a code symmetry group enables precise and efficient reasoning of code semantics. Our solution, SymC, develops a novel variant of self-attention that is provably equivariant to code symmetries from the permutation group defined over the program dependence graph. SymC obtains superior performance on five program analysis tasks, outperforming state-of-the-art code models, including GPT-4, without any pre-training. Our results suggest that code LLMs that encode the code structural prior via the code symmetry group generalize better and faster.
APA
Pei, K., Li, W., Jin, Q., Liu, S., Geng, S., Cavallaro, L., Yang, J. & Jana, S.. (2024). Exploiting Code Symmetries for Learning Program Semantics. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:40092-40113 Available from https://proceedings.mlr.press/v235/pei24b.html.

Related Material