{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/probabilistic-inference-in-language-models","title":"Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo","arxiv_id":"2404.17546","date":"2024-04-26","proceeding":null,"authors":["Stephen Zhao","Rob Brekelmans","Alireza Makhzani","Roger Grosse"],"abstract":"Numerous capability and safety techniques of Large Language Models (LLMs), including RLHF, automated red-teaming, prompt engineering, and infilling, can be cast as sampling from an unnormalized target distribution defined by a given reward or potential function over the full sequence. In this work, we leverage the rich toolkit of Sequential Monte Carlo (SMC) for these probabilistic inference problems. In particular, we use learned twist functions to estimate the expected future value of the potential at each timestep, which enables us to focus inference-time computation on promising partial sequences. We propose a novel contrastive method for learning the twist functions, and establish connections with the rich literature of soft reinforcement learning. As a complementary application of our twisted SMC framework, we present methods for evaluating the accuracy of language model inference techniques using novel bidirectional SMC bounds on the log partition function. These bounds can be used to estimate the KL divergence between the inference and target distributions in both directions. We apply our inference evaluation techniques to show that twisted SMC is effective for sampling undesirable outputs from a pretrained model (a useful component of harmlessness training and automated red-teaming), generating reviews with varied sentiment, and performing infilling tasks.","url_abs":"https://arxiv.org/abs/2404.17546v1","url_pdf":"https://arxiv.org/pdf/2404.17546v1.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"probabilistic-inference-in-language-models","repo_url":"https://github.com/silent-zebra/twisted-smc-lm","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok"}}],"tasks":[{"task_slug":"language-modelling","task_name":"Language Modelling"},{"task_slug":"prompt-engineering","task_name":"Prompt Engineering"},{"task_slug":"red-teaming","task_name":"Red Teaming"}],"methods":[{"method_slug":"focus","method_name":"Focus"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"syntology_url":"https://syntology.ai/paper/2404.17546","atlas_url":"https://app.syntology.ai/?focus=2404.17546","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2404.17546"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-25T09:33:49+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/silent-zebra/twisted-smc-lm","reach":{"status":"ok"}}],"summary":{"ran":11,"unverified":6},"by_repo_kind":{"official":{"samples":17,"ran":11,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":17,"samples":[{"code_sha256_prefix":"e5a65f3e5daae4a8","entry":"batch_reward_model","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"reward_models.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/reward_models.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"e5a65f3e5daae4a8"}},{"code_sha256_prefix":"1ad3697a8f2357c0","entry":"collect_info_across_trueposts","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"do_training_and_log_Z_bounds.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/do_training_and_log_Z_bounds.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"1ad3697a8f2357c0"}},{"code_sha256_prefix":"8df8935e525d9f54","entry":"get_all_seqs_up_to_output_len","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"custom_transformer_prob_utils.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/custom_transformer_prob_utils.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"8df8935e525d9f54"}},{"code_sha256_prefix":"7e325722c329023a","entry":"get_full_list_of_all_seqs_up_to_output_len","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"custom_transformer_prob_utils.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/custom_transformer_prob_utils.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"7e325722c329023a"}},{"code_sha256_prefix":"951de47945a1db46","entry":"get_new_params_twist_and_optim_twist_state","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"do_training_and_log_Z_bounds.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/do_training_and_log_Z_bounds.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"951de47945a1db46"}},{"code_sha256_prefix":"e2d64bdf2aa590e4","entry":"kl_div_jax","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"custom_transformer_prob_utils.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/custom_transformer_prob_utils.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"e2d64bdf2aa590e4"}},{"code_sha256_prefix":"f6fc75d8cc49afe3","entry":"kl_div_jax_sum_last_axis","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"backup/transformer_smc_huggingface.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/backup/transformer_smc_huggingface.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"f6fc75d8cc49afe3"}},{"code_sha256_prefix":"ad44f86f80299d7f","entry":"neg_beta_times_batch_reward_model_curry","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"reward_models.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/reward_models.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"ad44f86f80299d7f"}},{"code_sha256_prefix":"3260314e40b4993d","entry":"neg_beta_times_reward_model","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"reward_models.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/reward_models.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"3260314e40b4993d"}},{"code_sha256_prefix":"1121fae94eb827ac","entry":"populate_f_q_g_q_lists","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"get_kl_table.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/get_kl_table.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"1121fae94eb827ac"}},{"code_sha256_prefix":"f9edaa06c26334a2","entry":"stochastic_transformer_sample_iter","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"backup/transformer_smc_huggingface.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/backup/transformer_smc_huggingface.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"f9edaa06c26334a2"}},{"code_sha256_prefix":"9c7f287d70eb2ae3","entry":"get_logZ_midpoint_estimates","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"get_kl_table.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/get_kl_table.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"9c7f287d70eb2ae3"}},{"code_sha256_prefix":"d745ccf2ee440bf9","entry":"get_logZ_midpoint_to_use","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"get_kl_table.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/get_kl_table.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"d745ccf2ee440bf9"}},{"code_sha256_prefix":"96271cafe3d6e6b1","entry":"get_tokenizer","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"huggingface_models_custom.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/huggingface_models_custom.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"96271cafe3d6e6b1"}},{"code_sha256_prefix":"5a8ba055520815ed","entry":"inspect_and_record_evidence_setting_for_index","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"do_training_and_log_Z_bounds.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/do_training_and_log_Z_bounds.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"5a8ba055520815ed"}},{"code_sha256_prefix":"f7056ad862e2b7d9","entry":"load_ckpt","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"plot_bounds.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/plot_bounds.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"f7056ad862e2b7d9"}},{"code_sha256_prefix":"8c43715766ba7852","entry":"stochastic_transformer_sample","repo":"silent-zebra/twisted-smc-lm","repo_kind":"official","path":"backup/transformer_smc_huggingface.py","file_url":"https://github.com/silent-zebra/twisted-smc-lm/blob/HEAD/backup/transformer_smc_huggingface.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"8c43715766ba7852"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}