AlphaDev is an AI model based on the AlphaZero/MuZero Reinforcement Learning architecture. It's designed to optimize assembly code using a set of assembly instructions and a cost function which takes into account both correctness and performance.
pip install alphadev
import torch
from alpha_dev.model import AlphaDev
model = AlphaDev().cuda()
x = torch.randint(0, 256, (1, 1024)).cuda()
model(x) # (1, 1024, 20000)
AlphaDev consists of:
-
Representation Network:
f_rep
that outputs a latent representationht
of the stateSt
. -
Prediction Network:
f_pred
that predicts the expected return (the value)vˆt
and a policyπˆt
from a given latent state. -
Dynamics Network:
f_dyn
that predicts the next latent statehtk+1
and rewardrˆtk+1
resulting from a transition.
On reaching a new state, AlphaDev encodes the state into a latent representation using the representation network. The dynamics and prediction networks are used to simulate several trajectories that fill out a search tree by sampling state transitions.
The actions are selected using a strategy that balances exploration (trying new actions) and exploitation (progressing further down the subtree of the current best action).
Finally, the predicted policy is trained to match the visit counts of the MCTS policy in an attempt to distil the search procedure into a policy that will disregard nodes that are not promising.
AlphaDev, due to its general architecture, could potentially be adapted to solve a wide variety of optimization problems. Here are a few examples:
-
Route Optimization: For logistics companies, optimizing the routes of their fleet can result in significant cost savings. AlphaDev could be used to learn the optimal routes based on a variety of factors such as traffic, distance, and number of stops.
-
Job Scheduling: In computing, job scheduling is a key issue. AlphaDev could be used to learn the optimal schedule that maximizes the usage of computational resources and minimizes job completion time.
-
Stock Portfolio Optimization: AlphaDev could be used to learn the optimal mix of stocks to maximize return and minimize risk, given the current market conditions.
-
Game Playing: Similar to its ancestor AlphaZero, AlphaDev could potentially be used to master a wide variety of games, by learning the optimal strategies.
-
Drug Discovery: AlphaDev could be used to find the optimal chemical structure for a new drug that maximizes efficacy and minimizes side effects.
-
AssemblyGame This represents the Assembly Game RL environment. The state of the RL environment contains the current program and the state of memory and registers. Doing a step in this environment is equivalent to adding a new assembly instruction to the program (see the step method). The reward is a combination of correctness and latency reward after executing the assembly program over an input distribution. For simplicity of the overall algorithm we are not including the assembly runner, but assembly execution can be delegated to an external library (e.g. AsmJit).
-
AlphaDevConfig contains the main hyperparameters used for the AlphaDev agent. This includes configuration of AlphaZero, MCTS, and underlying networks.
-
play_game contains the logic to run an AlphaDev game. This include the MCTS procedure and the storage of the game.
-
RepresentationNet and PredictionNet contain the implementation the networks used in the AlphaZero algorithm. It uses a MultiQuery Transformer to represent assembly instruction
Future adaptations of AlphaDev could implement different learning algorithms or optimization techniques for specific domains or problem areas.
Please read CONTRIBUTING.md for details on our code of conduct, and the process for submitting pull requests to us.
This project is licensed under the MIT License - see the LICENSE.md file for details.
We appreciate the efforts of the researchers and developers who contributed to the development of the AlphaZero/MuZero architectures on which AlphaDev is based.
-
Add jax-based multi query attention:
MultiQueryAttentionBlock
-
add
ResBlockV2
-
add utils, terminal, is_correct, legal_actions