rStar-Math của Microsoft: Hướng dẫn kèm triển khai

Bài báo RStar-math của Microsoft giới thiệu một cách tiếp cận mới để giải bài toán toán học bằng cách kết hợp học tăng cường, suy luận ký hiệu và Monte Carlo Tree Search (MCTS).

Trong bài viết này, tôi sẽ khám phá khung RStar và các thành phần lõi của nó. Sau đó, tôi sẽ hướng dẫn bạn từng bước qua một bản triển khai đơn giản hóa minh họa các khái niệm chính bằng Gradio. Dù bản demo lấy cảm hứng từ bài báo, một số mức độ phức tạp đã được giản lược để dễ tiếp cận hơn.

rStar-Math của Microsoft là gì?

RStar-math hướng tới việc kết nối suy luận ký hiệu với khả năng khái quát hóa của các mô hình nơ-ron tiền huấn luyện. Khung này tích hợp các thành phần như Monte Carlo Tree Search (MCTS), mô hình ngôn ngữ tiền huấn luyện và học tăng cường để cho phép khám phá hiệu quả các chiến lược giải bài toán.

Ý tưởng cốt lõi là biểu diễn suy luận toán học như một quá trình tìm kiếm trên một cây có cấu trúc gồm các bước khả dĩ, nơi mỗi nút đại diện cho một nghiệm bộ phận hoặc trạng thái.

Nguồn: Guan et al., 2025

Một số lý do khiến rStar-Math đặc biệt thú vị với tôi là:

1. Bao gồm một mạng nơ-ron (mô hình chính sách) dự đoán hành động tiếp theo khi giải một bài toán toán học, định hướng cho việc khám phá của MCTS.
2. Một mạng (mô hình phần thưởng) đánh giá mức độ thành công của các hành động thực hiện trong quá trình rollout của MCTS và cung cấp phản hồi để huấn luyện.
3. RStar sử dụng các thư viện tính toán ký hiệu như SymPy cho các phép toán chính xác hoặc suy luận ký hiệu như giải phương trình hay tính đạo hàm.
4. Nhúng thuật toán Monte Carlo Tree Search để có hệ thống khám phá các lộ trình lời giải bằng mô phỏng, cân bằng giữa khám phá (thử đường đi mới) và khai thác.
5. Cơ chế phản hồi, trong đó các mô hình chính sách và phần thưởng được huấn luyện lặp lại dựa trên kết quả rollout MCTS, cải thiện dần quá trình ra quyết định.
6. Quá trình suy luận được cấu trúc như một cây phân cấp, nơi các nút là trạng thái và các cạnh là chuyển tiếp.

Tổng quan dự án demo: Trình giải bài toán với Gradio

Bản demo minh họa cách một mô hình chính sách và mô hình phần thưởng, kết hợp với suy luận ký hiệu dùng thư viện sympy , có thể xử lý bài toán toán học theo cách có cấu trúc. Các tính năng chính của bản triển khai này gồm:

1. Mô hình chính sách: Mạng nơ-ron dự đoán hành động tiếp theo trong quá trình giải bài toán.
2. Mô hình phần thưởng: Mạng đánh giá mức độ thành công của các hành động trong quá trình rollout MCTS.
3. Suy luận ký hiệu: Sử dụng SymPy cho tính toán toán học chính xác và giải phương trình.
4. Monte Carlo Tree Search: Triển khai phiên bản đơn giản của MCTS để khám phá lời giải hiệu quả.
5. Vòng lặp học tăng cường: Vòng huấn luyện cơ bản để cải thiện mô hình chính sách và phần thưởng dựa trên phản hồi.
6. Hỗ trợ phương trình một và nhiều biến: Người dùng có thể nhập một hoặc hai phương trình để tìm nghiệm cho các biến như x và y.

Để giữ bản demo đơn giản và tập trung, một số tính năng nâng cao được nêu trong bài báo nằm ngoài phạm vi của hướng dẫn này. Những tính năng đó gồm:

1. Khả năng mở rộng: Bài báo gốc dùng các mô hình tiền huấn luyện lớn và tài nguyên tính toán dồi dào. Bản demo dùng mạng nơ-ron nhỏ hơn và tránh tiền huấn luyện phức tạp.
2. Chiến lược MCTS nâng cao: Các kỹ thuật như UCT thích ứng và chiến lược khám phá đa dạng chưa được triển khai đầy đủ.
3. Khái quát hóa tác vụ: Bản triển khai chỉ tập trung vào giải phương trình đại số, trong khi RStar được thiết kế để khái quát trên dải tác vụ toán học rộng hơn.
4. Tập dữ liệu: Thay vì dùng tập dữ liệu chọn lọc để huấn luyện, bản demo dựa vào suy luận ký hiệu và dữ liệu người dùng cung cấp.

Bước 1: Yêu cầu tiên quyết

Bản demo được chia thành nhiều thành phần, mỗi phần phản ánh một phần của phương pháp RStar. Trước khi bắt đầu, hãy đảm bảo bạn đã cài đặt:

• Python 3.8+
• Các thư viện bắt buộc: Cài các gói Python cần thiết bằng pip:

pip install requests gradio, sympy Được hỗ trợ bởi

Sau đó, import các thư viện này:

import gradio as gr
import numpy as np
import torch
import re
import torch.nn as nn
import torch.optim as optim
from sympy import symbols, Eq, solve, N, sin, cos, tan, exp, log, E, sympify
from random import choiceĐược hỗ trợ bởi

Khi đã cài đặt xong mọi phụ thuộc, hãy thiết lập các thành phần chính.

Bước 2: Mạng nơ-ron cho Chính sách và Phần thưởng

Các mạng này là phiên bản gọn nhẹ của các mô hình mô tả trong bài báo, dùng để dự đoán hành động tiếp theo và đánh giá mức độ thành công. Mô hình chính sách dự đoán các bước tiếp theo để giải phương trình đã cho. Nó dùng mạng nơ-ron truyền thẳng để xử lý biểu diễn đã mã hóa của bài toán.

Tương tự, mô hình phần thưởng đánh giá nghiệm bộ phận để định hướng quá trình MCTS. Cả hai mô hình đều được hiện thực bằng PyTorch.

class PolicyModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x
class RewardModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return xĐược hỗ trợ bởi

Tiếp theo, ta xây dựng lớp nút cho cây MCTS.

Bước 3: Lớp TreeNode để biểu diễn trạng thái MCTS

Lớp TreeNode biểu diễn các nút trong cây MCTS. Mỗi nút tương ứng với một trạng thái trong quá trình tìm kiếm, bao gồm:

• Trạng thái (ví dụ: các phương trình hoặc nghiệm bộ phận).
• Tham chiếu tới nút cha.
• Danh sách các nút con (các trạng thái được mở rộng).
• Số lượt thăm và giá trị Q, theo dõi tần suất khám phá nút và tổng phần thưởng tích lũy.

class TreeNode:
    """Represents a node in the MCTS tree."""
    def __init__(self, state, parent=None):
        self.state = state  # Current state
        self.parent = parent
        self.children = []
        self.visits = 0
        self.q_value = 0.0  # Accumulated rewards
    def is_fully_expanded(self):
        return len(self.children) > 0
    def best_child(self, exploration_weight=1.4):
        """Select the best child using UCT formula."""
        def uct_value(child):
            return (child.q_value / (child.visits + 1e-6)) + exploration_weight * np.sqrt(np.log(self.visits + 1) / (child.visits + 1e-6))
        return max(self.children, key=uct_value)
    def add_child(self, child_state):
        """Add a child node with the given state."""
        child = TreeNode(state=child_state, parent=self)
        self.children.append(child)
        return childĐược hỗ trợ bởi

Khi đã có cấu trúc cơ bản, ta sẽ làm việc với các thành phần lõi tiếp theo của demo.

Bước 4: Lớp MathSolver

Lớp MathSolver là lõi của bản demo, kết hợp suy luận ký hiệu với tìm kiếm có hướng dẫn bởi mạng nơ-ron. Nó hiện thực một số thành phần chính:

A. Mô hình chính sách và phần thưởng

PolicyModel dự đoán các bước tiếp theo để giải phương trình, trong khi RewardModel đánh giá mức độ thành công của các nghiệm bộ phận hoặc hoàn chỉnh.

class MathSolver:
    def __init__(self, dataset=None):
        self.dataset = dataset or []  # Dataset of math problems
        self.policy_model = PolicyModel(input_size=128, hidden_size=64, output_size=4)  
        self.reward_model = RewardModel(input_size=128, hidden_size=64, output_size=1)  
        self.policy_optimizer = optim.Adam(self.policy_model.parameters(), lr=0.001)
        self.reward_optimizer = optim.Adam(self.reward_model.parameters(), lr=0.001)
        self.execution_context = {}  Được hỗ trợ bởi

Phương thức trên khởi tạo lớp MathSolver bằng cách thiết lập các thành phần cần thiết để giải bài toán. Nó có thể nhận tùy chọn một tập dữ liệu bài toán và khởi tạo hai mạng nơ-ron: mô hình chính sách dự đoán hành động tiếp theo, và mô hình phần thưởng đánh giá mức độ thành công của các hành động. 

Giờ ta đã có hàm chính sách và phần thưởng. Tiếp theo, ta cần phân tích cú pháp và mã hóa các phương trình đầu vào.

B. Phân tích và mã hóa phương trình

Các phương trình được phân tích bằng sympy và mã hóa thành các vector đặc trưng để mô hình chính sách và phần thưởng xử lý.

def encode_problem(self, problem):
 # Advanced encoding using symbolic representation and problem length
     variables = len(re.findall(r'[a-zA-Z]', problem))
     operators = len(re.findall(r'[\+\-\*/\^]', problem))
      problem_length = len(problem)
      return np.array([variables, operators, problem_length] + [0] * 125)Được hỗ trợ bởi

Phương thức encode_problem chuyển một bài toán toán học thành biểu diễn số có kích thước cố định cho các mô hình. Nó trích xuất các đặc trưng như số lượng biến, toán tử và độ dài bài toán, mã hóa thành mảng NumPy 128 chiều. Biểu diễn này nắm bắt cấu trúc bài toán, giúp mô hình xử lý hiệu quả.

C. Dự đoán của mô hình chính sách

Đoạn mã sau sinh các bước để giải các phương trình đã cho, bao gồm định nghĩa biến, tạo phương trình và giải chúng.

def policy_model_predict(self, equation1, equation2=None):
    try:
        equations = []
        if equation1:
            equations.append(sympify(equation1.strip()))  # Sympify only equations
        if equation2:
            equations.append(sympify(equation2.strip()))
        all_variables = set()
        for eq in equations:
            all_variables.update(eq.free_symbols)
        var_definitions = [f"{v} = symbols('{v}')" for v in all_variables]
        steps = [
            ("Define variables", "\n".join(var_definitions)),
            ("Define equation(s)", f"equations = {equations}"),
            ("Solve equation(s)", f"solution = solve(equations, {list(all_variables)})"),
            ("Print solution", "print(solution)")
        ]
        return steps
    except Exception as e:
        print(f"Error during policy model prediction: {e}")
        return []Được hỗ trợ bởi

Hàm policy_model_predict phân tích các phương trình đầu vào bằng sympify của SymPy để đảm bảo chúng là biểu thức toán hợp lệ. Sau đó, hàm xác định tất cả biến xuất hiện trong các phương trình và giải chúng bằng hàm solve của SymPy. Phương thức này đóng vai trò như một hướng dẫn cấp cao cho quy trình giải bài toán. 

D. Dự đoán của mô hình phần thưởng

Phương thức reward_model_predict đóng vai trò quan trọng trong học tăng cường bằng cách cung cấp phản hồi cho các hành động thực hiện trong các rollout của MCTS. 

  
  def reward_model_predict(self, steps, success):
        encoded_steps = self.encode_problem(str(steps))
        encoded_steps = torch.tensor(encoded_steps, dtype=torch.float32)
        reward = self.reward_model(encoded_steps)
        return reward.item() if success else -reward.item()Được hỗ trợ bởi

Hàm này mã hóa các bước giải thành dạng số và đánh giá qua mô hình phần thưởng, trả về phần thưởng dương khi thành công và âm khi thất bại. Phản hồi này huấn luyện mô hình chính sách, giúp ưu tiên các hành động hiệu quả và cải thiện quyết định. Với các hàm dự đoán của mô hình chính sách và phần thưởng đã sẵn sàng, giờ ta có thể làm việc với tác vụ thực thi.

E. Hàm thực thi mã

Phương thức này xử lý nghiệm nhiều biến dưới dạng bộ hoặc từ điển và chuyển kết quả ký hiệu sang xấp xỉ số bằng hàm N của SymPy. 

    def execute_code(self, code):
        try:
            # Ensure necessary imports and variables are in the execution context
            exec("from sympy import symbols, Eq, solve, N, sin, cos, tan, exp, log, E", self.execution_context)
            # Dynamically initialize variables in the context
            for var_def in self.execution_context.get('var_definitions', []):
                exec(var_def, self.execution_context)
            exec(code, self.execution_context)
            if "solution" in self.execution_context:
                symbolic_solution = self.execution_context["solution"]
                # Handle multi-variable solutions as tuples
                if isinstance(symbolic_solution, list):
                    self.execution_context["solution"] = [tuple(map(N, sol)) if isinstance(sol, tuple) else N(sol) for sol in symbolic_solution]
                elif isinstance(symbolic_solution, dict):
                    self.execution_context["solution"] = {k: N(v) for k, v in symbolic_solution.items()}
                else:
                    self.execution_context["solution"] = N(symbolic_solution)
            return True
        except Exception as e:
            print(f"Error executing code: {e}")
            return FalseĐược hỗ trợ bởi

Phương thức này đảm bảo tính toán chính xác và hỗ trợ linh hoạt nhiều định dạng nghiệm. Nếu có lỗi, nó ghi lại vấn đề và trả về False, giúp xử lý lỗi vững chắc.

F. Monte Carlo Tree Search (MCTS)

Phương thức MCTS lặp lại việc chọn các trạng thái tốt nhất, mở rộng cây tìm kiếm và mô phỏng các lời giải khả dĩ. Phần thưởng từ mô phỏng được lan truyền ngược để cải thiện quyết định.

def mcts(self, equation1, equation2=None, num_rollouts=10):
    root = TreeNode(state=(equation1, equation2))
    for _ in range(num_rollouts):
        # Selection
        node = root
        while node.is_fully_expanded() and node.children:
            node = node.best_child()
        # Expansion
        if not node.is_fully_expanded():
            steps = self.policy_model_predict(*node.state)
            for step, code in steps:
                child_state = (step, code)
                node.add_child(child_state)
        # Simulation
        success = True
        for step, code in steps:
            if not self.execute_code(code):
                success = False
                break
        # Backpropagation
        reward = self.reward_model_predict(steps, success)
        while node:
            node.visits += 1
            node.q_value += reward
            node = node.parent
    return root.best_child().state if root.children else NoneĐược hỗ trợ bởi

Phương thức mcts thực hiện tuần tự bốn bước chính: 

• Chọn (Selection): Di chuyển xuống cây theo các nút con tốt nhất.
• Mở rộng (Expansion): Tạo các nút con mới bằng phương thức policy_model_predict.
• Mô phỏng (Simulation): Thực thi tất cả hành động và xác định thành công.
• Lan truyền ngược (Backpropagation): Tính phần thưởng bằng phương thức reward_model_predict và lan truyền lên cây để cập nhật giá trị các nút. 

Sau số lượt rollout nhất định, phương thức trả về trạng thái của nút con tốt nhất, đại diện cho lời giải hứa hẹn nhất trong quá trình tìm kiếm.

G. Thực thi lời giải

Phương thức solve điều phối toàn bộ quy trình, từ phân tích phương trình đến thực thi và kiểm chứng lời giải.

def solve(self, equation1, equation2=None):
    self.execution_context = {}
    steps = self.policy_model_predict(equation1, equation2)
    variables = set()
    for eq in [equation1, equation2] if equation2 else [equation1]:
        if eq:
            variables.update(sympify(eq.strip()).free_symbols)
    self.execution_context['var_definitions'] = [f"{v} = symbols('{v}')" for v in variables]
    steps_output = ["Best solution found:"]
    for step, code in steps:
        steps_output.append(f"Step: {step}")
        steps_output.append(f"Code: {code}")
        if self.execute_code(code):
            steps_output.append("Execution successful.")
        else:
            steps_output.append("Execution failed.")
    if "solution" in self.execution_context:
        final_answer = self.execution_context["solution"]
        if isinstance(final_answer, dict):
            for var, value in final_answer.items():
                steps_output.append(f"{var} = {value}")
        elif isinstance(final_answer, list):
            for solution in final_answer:
                if isinstance(solution, tuple):
                    for idx, var in enumerate(variables):
                        steps_output.append(f"{list(variables)[idx]} = {solution[idx]}")
                else:
                    steps_output.append(f"Solution: {solution}")
        else:
            steps_output.append(f"Final Answer: {final_answer}")
    else:
        steps_output.append("No final answer found.")
    return "\n".join(steps_output)Được hỗ trợ bởi

Phương thức solve xử lý một hoặc hai phương trình do người dùng cung cấp bằng cách khởi tạo ngữ cảnh thực thi và tạo các bước qua policy_model_predict. Nó thực thi từng bước, ghi nhận tiến độ và báo cáo thành công hay thất bại. Lời giải, bao gồm nghiệm đơn biến và đa biến, được định dạng kèm tên biến và giá trị để rõ ràng. Nếu không tìm thấy lời giải, sẽ hiển thị thông báo phù hợp.

Ta đã có đầy đủ các thành phần lõi, giờ có thể xây dựng ứng dụng Gradio.

Bước 5: Tạo giao diện thân thiện với Gradio

Giao diện Gradio cho phép người dùng nhập phương trình (một hoặc nhiều), giải và xem kết quả tương tác.

with gr.Blocks() as app:
    gr.Markdown("# Math Problem Solver with Advanced Multi-Step Reasoning and Learning")
    with gr.Row():
        equation1_input = gr.Textbox(label="Enter the first equation (e.g., x + y - 3)", placeholder="x + y - 3")
        equation2_input = gr.Textbox(label="Enter the second equation (optional, e.g., x - y - 1)", placeholder="x - y - 1")
    solve_button = gr.Button("Solve")
    solution_output = gr.Textbox(label="Solution", interactive=False)
    solve_button.click(solve_math_problem, inputs=[equation1_input, equation2_input], outputs=[solution_output])
app.launch(debug=True)Được hỗ trợ bởi

Đoạn mã trên tạo một giao diện người dùng Gradio để giải phương trình toán với suy luận nâng cao. Giao diện được bao trong container gr.Blocks, gồm hai trường nhập bằng gr.Textbox: một cho phương trình thứ nhất (bắt buộc) và một cho phương trình thứ hai (tùy chọn).

Kết quả hiển thị trong một gr.Textbox có nhãn "Solution". Lệnh interface.launch() khởi chạy ứng dụng Gradio trên trình duyệt và cờ debug=True bật nhật ký chi tiết để hỗ trợ xử lý lỗi.

Bước 6: Kiểm thử và xác thực

Đến lúc kiểm thử ứng dụng Trình giải bài toán. Dưới đây là một số bài thử tôi đã chạy:

1. Một phương trình một biến: Tôi thử tìm các giá trị khả dĩ của một biến x với đầu vào là một phương trình.

2. Bài toán nhiều phương trình nhiều biến: Tôi đưa vào hai phương trình với hai biến để tìm các giá trị khả dĩ của các biến x và y.

Khả năng mở rộng thêm

 Bản demo này là phiên bản cơ bản của những gì ta có thể đạt được với phương pháp rStar-math. Vẫn còn nhiều việc có thể làm để mở rộng khả năng của nó. 

• Bổ sung các mô hình ngôn ngữ tiền huấn luyện để tăng cường khả năng suy luận của mô hình chính sách.
• Tích hợp các chiến lược MCTS nâng cao hoặc thử kỹ thuật khám phá đa dạng để cải thiện hiệu quả và độ chính xác của quá trình tìm kiếm.
• Mở rộng demo để giải các phương trình đa thức bậc cao và hệ phương trình phức tạp.
• Huấn luyện mô hình trên tập dữ liệu phương trình lớn hơn để cải thiện khả năng khái quát.
• Mở rộng demo để xử lý các tác vụ suy luận ngoài lĩnh vực toán.

Bạn có thể tham khảo kho mã gốc của bài báo rStar-math trên GitHub.

Kết luận

Bản demo này trình diễn một triển khai thực tiễn của suy luận nhiều bước để giải phương trình toán học. Bằng cách kết hợp mạng nơ-ron, suy luận ký hiệu và MCTS, nó cho thấy cách các kỹ thuật AI tiên tiến có thể xử lý các tác vụ suy luận có cấu trúc. Những nâng cấp trong tương lai có thể đưa nó tiến gần hơn tới đầy đủ khả năng của khung RStar.