PRefLexOR: Preference-based Recursive Language Modeling for Exploratory Optimization of Reasoning and Agentic Thinking

We introduce PRefLexOR (Preference-based Recursive Language Modeling for Exploratory Optimization of Reasoning), a framework that combines preference optimization with concepts from Reinforcement Learning (RL) to enable models to self-teach through iterative reasoning improvements. Central to PRefLexOR are thinking tokens, which explicitly mark reflective reasoning phases within model outputs, allowing the model to recursively engage in multi-step reasoning, revisiting, and refining intermediate steps before producing a final output. The foundation of PRefLexOR lies in Odds Ratio Preference Optimization (ORPO), where the model learns to align its reasoning with human-preferred decision paths by optimizing the log odds between preferred and non-preferred responses. The integration of Direct Preference Optimization (DPO) further enhances model performance by using rejection sampling to fine-tune reasoning quality, ensuring nuanced preference alignment. This hybrid approach between ORPO and DPO mirrors key aspects of RL, where the model is continuously guided by feedback to improve decision-making and reasoning. Active learning mechanisms allow PRefLexOR to dynamically generate new tasks, reasoning steps, and rejected answers on-the-fly during training. This adaptive process enables the model to self-teach as it continually improves through real-time feedback and recursive processing.

Our method diverges from traditional approaches by not relying on pre-generated datasets; instead, it dynamically generates new tasks, reasoning steps, and feedback on the fly, allowing the model to continuously adapt and improve in real time. Recursive optimization within the thinking token framework introduces iterative feedback loops, where the model refines its reasoning, much like policy refinement in RL, achieving deeper coherence, consistency, and adaptability. By recursively optimizing reasoning through feedback-driven learning, PRefLexOR achieves significant flexibility in its ability to handle complex tasks, learning and evolving its cognitive abilities autonomously. This framework advances the field of cognitive alignment by demonstrating that models can iteratively teach themselves to reason with greater depth and reflectivity, akin to an RL-based self-improving system capable of solving open-domain problems with superior reasoning depth and logic. Our implementation is straightforward and can be Incorporated into any existing pretrained LLM. The approach is demonstrated in use cases of materials design applications, where a small language model is trained to develop sophisticated reasoning capabilities. Thereby, PRefLexOR builds a dynamic knowledge graph by generating questions from random text and using Retrieval-Augmented Generation (RAG) to retrieve contextually relevant data from the entire corpus, facilitating recursive reasoning through complex interactions between similar nodes in the embedding space.

Source code: https://github.com/lamm-mit/PRefLexOR

Figure 1: Illustration of the workflow and design principles behind generative materials informatics. Panel a: The process of transforming information into knowledge and actionable outcomes. Each individual piece of information (left) is synthesized into a network of interconnected knowledge, leading to informed decisions and innovative designs (right). Panel b: Conventional approaches in materials science rely on data-driven models, partial differential equations (PDEs), and experimental results, focusing on single-step predictions. Panel c: In contrast, generative materials informatics models built on the PRefLexOR framework proposed in this paper use 'thinking' and 'reflection' explicitly by incorporating iterative reasoning and contextual understanding, allowing for more complex, multi-step predictions. This approach expands from single inference steps, includes multiple modalities of data and responses, integrates real-world feedback and physics, and leverages self-assessment and self-learning. Using using reinforcement learning (RL) principles, the discovery of principles or the solution of specific tasks is further inspired by biological paradigms, using bio-inspired neural network designs. These advanced methods support continuous improvement in material predictions, enabling more adaptable and intelligent designs

Figure 2: PRefLexOR Recursive Reasoning Algorithm: An iterative approach leveraging a fine-tuned Reasoning Model and a general-purpose Critic Model to generate, refine, and optionally integrate responses. The process involves generating initial responses, extracting reflections, improving thinking processes, and creating new responses based on refined thinking, with an optional final integration step. The algorithm relies on extracting thinking processes (indicated via <|thinking|>...<|/thinking|>) and reflection processes (indicated via <|reflect|>...<|/reflect|>). The use of special tokens allows us to easily construct such agentic modeling as it facilitates pausing inference, improving the strategy, and re-generating improved answers. The sampled responses can either be used in their final state or integrated into an amalgamated response that shows very rich facets in the scientific process.

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name='lamm-mit/PRefLexOR_ORPO_DPO_EXO_10242024'
model = AutoModelForCausalLM.from_pretrained(model_name,     
    torch_dtype =torch.bfloat16,
    attn_implementation="flash_attention_2",device_map="auto",trust_remote_code=True,
    )
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True,
                                          use_fast=False,
                                         )

The model produces a structured thinking section before answering, indicated by these special tokens:

thinking_start = '<|thinking|>'
thinking_end = '<|/thinking|>'

Inference examples

We provide several examples for how to use this model, using Colab notebooks and Hugging Face transformers.

Colab PRefLexOR Inference: Thinking and Agentic Reflection

Simple inference:

from PRefLexOR import *

txt = 'What is the relationship between materials and music? Brief answer.' + f' Use {think_start}.'

output_text, messages = generate_local_model(
    model=model, 
    tokenizer=tokenizer, 
    prompt=txt, 
    system_prompt='',  
    num_return_sequences=1, 
    repetition_penalty=1.0, 
    temperature=0.1, 
    max_new_tokens=2024, 
    messages=[], 
    do_sample=True
)

print(output_text)

Extract thinking and output:

thinking    = extract_text(output_text, thinking_start=think_start, thinking_end=think_end)[0].strip()
answer_only = extract_text(output_text, thinking_start=think_end, thinking_end="NONE").strip()
print ("THINKING:\n\n", thinking)
print ("ANSWER:\n\n", answer_only)

Recursive inference usingh multi-agentic modeling

from transformers import AutoModelForCausalLM, AutoTokenizer

# Load reasoning model
model_name='lamm-mit/PRefLexOR_ORPO_DPO_EXO_10242024'
model = AutoModelForCausalLM.from_pretrained(model_name,     
    torch_dtype =torch.bfloat16,
    attn_implementation="flash_attention_2",device_map="auto",trust_remote_code=True,
    )
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True,
                                          use_fast=False,
                                         )
# Load critic model
model_name_critic = "meta-llama/Llama-3.2-3B-Instruct"

critic_model = AutoModelForCausalLM.from_pretrained(
    model_name_critic, 
    torch_dtype=torch.bfloat16, 
    attn_implementation="flash_attention_2", 
    device_map="auto", 
    trust_remote_code=True
)

Example inference

output_text, output_list, output_text_integrated = recursive_response_from_thinking(
    model=model, 
    tokenizer=tokenizer, 
    model_critic=critic_model, 
    tokenizer_critic=tokenizer,  #same tokenizer in our case
    question="Develop an idea of how graphene can be combined with silk fibers to create a filtration membrane.", 
    N=3, 
    temperature=0.1, 
    temperature_improvement=0.1, 
    system_prompt="You are a helpful assistant.", 
    system_prompt_critic="You carefully improve responses, with attention to detail, and following all directions.",
    verbatim=False,

Printing the output:

for i, item in enumerate(output_list):
    print (f"i={i}", 64*"-")
    print (item)

print (64*"#")
print ("INTEGRATED RESPONSE:")
print (output_text_integrated)
print (64*"#")

Example integrated response:

Developing an idea of how graphene can be combined with silk fibers to create a filtration membrane involves several steps and considerations. The process begins with the functionalization of graphene oxide (GO) to improve its interaction with silk fibers. This can be achieved by treating GO with chemicals like sodium hydroxide (NaOH) or sodium borohydride (NaBH4) to introduce hydroxyl (-OH) or boronate (-B) groups, respectively.

Next, silk fibers are prepared by dissolving them in a solvent like N,N-dimethylformamide (DMF) or N,N-dimethylacetamide (DMAc). The silk solution is then electrospun to create fibers with a desired diameter and alignment. To create the graphene-silk hybrid membrane, the GO-functionalized silk fibers are mixed with GO powder in a solvent like DMF or DMAc. The mixture is then electrospun again to create the hybrid membrane.

To achieve desired mechanical properties, the spacing between GO layers can be controlled. This can be done by adjusting the concentration of GO powder, the solvent used, or the electrospinning parameters like voltage and flow rate. Silk fibers can be oriented during the electrospinning process to enhance filtration efficiency. This can be achieved by using a rotating drum or a vibrating plate to align the fibers.

The graphene-silk hybrid membrane's filtration performance can be evaluated by measuring its pore size, water flux, and pressure drop. The pore size can be controlled by adjusting the GO concentration or the electrospinning parameters. The water flux can be improved by increasing the silk fiber concentration or using a higher voltage during electrospinning.

The mechanical properties of the graphene-silk hybrid membrane can be assessed by measuring its tensile strength and Young's modulus. The tensile strength can be improved by increasing the GO concentration or using a higher voltage during electrospinning. The Young's modulus can be controlled by adjusting the silk fiber concentration or the GO concentration.

The graphene-silk hybrid membrane can be used for various applications, including water filtration, wound dressing, and even tissue engineering. Its unique properties, such as high surface area, good electrical conductivity, and mechanical strength, make it suitable for these applications.

A simple representation of the process can be shown as follows:

In this representation, DS3 refers to the silk fiber used for the hybrid membrane, and GO refers to the graphene oxide used for functionalization and hybridization. The process involves functionalization, electrospinning, hybridization, and orientation of silk fibers to create the graphene-silk hybrid membrane. The resulting membrane can be used for filtration and other applications.

The graphene-silk hybrid membrane offers several advantages, including enhanced mechanical properties due to the combination of GO and silk fibers, improved filtration efficiency due to the orientation of silk fibers, scalability for industrial applications due to the electrospinning technique, and potential for addressing specific contaminants or particles due to the unique properties of GO and silk fibers.

However, there are also potential challenges associated with the graphene-silk hybrid membrane, such as achieving consistent and controlled GO interlayer spacing to balance mechanical properties and filtration performance, ensuring the orientation of silk fibers does not compromise the mechanical properties of the membrane, scaling up the production of the hybrid membrane for industrial applications while maintaining its filtration performance and mechanical properties.

To address these challenges, it is essential to develop a comprehensive understanding of the material properties and processing conditions that affect the graphene-silk hybrid membrane. This can be achieved through a combination of theoretical modeling, experimental characterization, and simulation studies.

In conclusion, combining graphene with silk fibers to create a filtration membrane is a promising approach that offers several advantages and potential applications. However, it requires careful consideration of the material properties and processing conditions to achieve the desired mechanical properties and filtration performance. By understanding the underlying principles and challenges associated with this approach, researchers and engineers can develop innovative solutions for various applications, including water filtration, wound dressing, and tissue engineering.

Citation

@article{buehler2024PRefLexOR,
      title={PRefLexOR: Preference-based Recursive Language Modeling for Exploratory Optimization of Reasoning and Agentic Thinking}, 
      author={Markus J. Buehler},
      year={2024},
      eprint={2410.12375},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2410.12375}, 
}

lamm-mit
/

PRefLexOR_ORPO_DPO_EXO_10242024