Learning Iterative Reasoning through Energy Diffusion

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Du, Yilun, Mao, Jiayuan, Tenenbaum, Joshua B.
Natura: Preprint
Pubblicazione: 2024
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866909225263824896
author Du, Yilun
Mao, Jiayuan
Tenenbaum, Joshua B.
author_facet Du, Yilun
Mao, Jiayuan
Tenenbaum, Joshua B.
contents We introduce iterative reasoning through energy diffusion (IRED), a novel framework for learning to reason for a variety of tasks by formulating reasoning and decision-making problems with energy-based optimization. IRED learns energy functions to represent the constraints between input conditions and desired outputs. After training, IRED adapts the number of optimization steps during inference based on problem difficulty, enabling it to solve problems outside its training distribution -- such as more complex Sudoku puzzles, matrix completion with large value magnitudes, and pathfinding in larger graphs. Key to our method's success is two novel techniques: learning a sequence of annealed energy landscapes for easier inference and a combination of score function and energy landscape supervision for faster and more stable training. Our experiments show that IRED outperforms existing methods in continuous-space reasoning, discrete-space reasoning, and planning tasks, particularly in more challenging scenarios. Code and visualizations at https://energy-based-model.github.io/ired/
format Preprint
id arxiv_https___arxiv_org_abs_2406_11179
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Learning Iterative Reasoning through Energy Diffusion
Du, Yilun
Mao, Jiayuan
Tenenbaum, Joshua B.
Machine Learning
Artificial Intelligence
We introduce iterative reasoning through energy diffusion (IRED), a novel framework for learning to reason for a variety of tasks by formulating reasoning and decision-making problems with energy-based optimization. IRED learns energy functions to represent the constraints between input conditions and desired outputs. After training, IRED adapts the number of optimization steps during inference based on problem difficulty, enabling it to solve problems outside its training distribution -- such as more complex Sudoku puzzles, matrix completion with large value magnitudes, and pathfinding in larger graphs. Key to our method's success is two novel techniques: learning a sequence of annealed energy landscapes for easier inference and a combination of score function and energy landscape supervision for faster and more stable training. Our experiments show that IRED outperforms existing methods in continuous-space reasoning, discrete-space reasoning, and planning tasks, particularly in more challenging scenarios. Code and visualizations at https://energy-based-model.github.io/ired/
title Learning Iterative Reasoning through Energy Diffusion
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2406.11179