multi-head design<\/em>.<\/li>\n<\/ul>\nWe’ll save the part for multi-head design for last. For now, let’s focus on the essential components to lay the groundwork for multi-head design use cases.<\/p>\n
Designing a deep neural network in PyTorch can be as simple as using a pre-defined model class like Linear as such: layer = Linear(3, 5)<\/code> where 3<\/code> is the number of input features and 5<\/code> is the number of output features (a.k.a.\u00a0nodes\u00a0<\/em>in many visualizations.) This syntax doesn’t account for the concept of activation functions, however, but PyTorch provides support there too.<\/p>\nThis can be implemented simply as relu(layer(x))<\/code> where relu<\/code> is an activation function from PyTorch’s functional<\/code> Modele, layer<\/code> is the Linear model we just created, and x<\/code> are the input parameters. Here, the relu<\/code> function\u00a0transforms\u00a0<\/em>the output of the network layer. Here’s an example of defining each layer separately:<\/p>\nimport torch\r\nimport torch.nn as nn\r\nimport torch.nn.functional as F\r\n\r\n# verbose definition of params\r\ninput_features = 5\r\noutput_features = 2\r\nhidden_layer_size = 128\r\n\r\n# section 1: network design\r\nlayer_1 = nn.Linear(input_features, hidden_layer_size)\r\nlayer_2 = nn.Linear(hidden_layer_size, hidden_layer_size)\r\nlayer_3 = nn.Linear(hidden_layer_size, output_features)\r\n\r\n# section 2: forward pass\r\nx = ...\r\nx = F.relu(layer_1(x))\r\nx = F.relu(layer_2(x))\r\nx = layer_3(x) # <----- the final output layer\r\n\r\n# an example of applying the action\r\naction = torch.argmax(x).item()<\/pre>\nHere we see a simple three-layer design:<\/p>\n
\n- layer_1<\/strong>: the layer that takes the initial inputs and outputs into the hidden layer.<\/li>\n
- layer_2<\/strong>: the “hidden layer” that makes this a “deep” neural network.<\/li>\n
- layer_3<\/strong>: the output layer, that translates the network learning into a probability space for each action. For example a Buy\/Sell signal.<\/li>\n<\/ul>\n
In the block labeled “section 2” we see an implementation of a “forward pass” via which input passes<\/em> through the network, being transformed by weights, biases, and activation functions (relu in this case) into the output layer. Practically, section2<\/code> would be implemented as the forward<\/code> method in a subclass of the nn.Module<\/code> class in PyTorch. Here, however, it serves to illustrate a basic use only.<\/p>\nMulti-Headed Deep Neural Network in PyTorch<\/h2>\n
Let’s consider what a multi-headed design might look like — then we’ll move on to discussing why it might be beneficial:<\/p>\n
import torch\r\nimport torch.nn as nn\r\nimport torch.nn.functional as F\r\n\r\n# defines out network params\r\ninput_features = 5\r\noutput_features = 2\r\nhidden_layer_size = 128\r\n\r\n# defines the network structure\r\nlayer_1 = nn.Linear(input_features, hidden_layer_size)\r\nlayer_2 = nn.Linear(hidden_layer_size, hidden_layer_size)\r\n\r\n# the multi-headed components, isolated from each other\r\ntrade_action = nn.Linear(hidden_layer_size, output_features)\r\nleverage_amt = nn.Linear(hidden_layer_size, out_features=100)\r\nrisk_amt = nn.Linear(hidden_layer_size, out_features=100)\r\n\r\n# example of a forward pass\r\nx = ...\r\nx = F.relu(layer_1(x))\r\nx = F.relu(layer_2(x))\r\n\r\n# separate actions for each \"head\"\r\nbuy_or_sell = torch.argmax(trade_action(x)).item\r\nleverage = torch.argmax(use_leverage(x)).item\r\nrisk = torch.argmax(risk_amt(x)).item\r\n<\/pre>\nHere we see a stark difference in our forward-pass approach. Rather than outputting a single value for the entire network, we implement three separate network “heads” for three conceptually different concerns:<\/p>\n
\n- trade_action<\/strong>: The signal to Buy\/Sell an asset.<\/li>\n
- leverage_amt<\/strong>: The amount of leverage to use.<\/li>\n
- risk_amt<\/strong>: The amount of risk to take.<\/li>\n<\/ul>\n
Here we see that all the input values into the network first pass through the input layer and then the hidden layer. After that, however, rather than all actions being encoded into a single layer, they are isolated. Two immediate benefits to note are:<\/p>\n
\n- Conceptual separation of tasks\/outputs.<\/li>\n
- Parameter separation of learning.<\/li>\n<\/ul>\n
The first benefit helps keep a more semantic flow during development and allows for stronger alignment with domain-specific task definitions. The second benefit is more foundational and tough to objectively measure. Does it make sense to separate these actions? There’s no universal answer to this question and each network designer should carefully consider the problem domain, what’s being optimized for, and the developer workflow. This particular case, an example of how this might make sense:<\/p>\n
# reset a training environment\r\nagent = ...\r\nenv = ...\r\nstate, info = env. reset()\r\ndone = False\r\nwhile not done:\r\n \r\n # generate an action(s) from the network\r\n buy_sell, leverage_amount, risk_amount = agent(state)\r\n\r\n # apply the action to the environment\r\n state, reward, terminal, info = env(buy_sell, leverage_amount, risk_amount)\r\n<\/pre>\nHere we note an immediate semantic benefit of having three distinct components of an agent’s action returned from the network. This can help avoid simple syntactic errors during the development of training routines and represent a clearer intent in code. In this case, these actions are doing the following:<\/p>\n
\n- buy_sell<\/strong>: Whether to Buy or Sell an asset.<\/li>\n
- leverage_amount<\/strong>: The amount of leverage (if relevant) to use on the order.<\/li>\n
- risk_amount<\/strong>: The risk to accept during order sizing.<\/li>\n<\/ul>\n
How these action components are applied in the environment is case-specific. However, the clarity of their intent and simplified application process are broadly applicable and good software design practices<\/a> in my opinion.<\/p>\nDiscussion<\/h2>\n
Multi-headed network designs offer novel discretization of network components. This offers both semantic and performance improvements — depending on the use case. Modern LLMs tend to use multi-head components directing their “attention heads”[2] for previously unmatched emergent properties.<\/p>\n
Essentially, the multi-headed components of LLM models are largely attributed as having given rise to the new generation of models, such as ChatGPT, that offer significant performance improvements.<\/p>\n
Will dropping in a multi-head design to your network turn your startup into the next OpenAI — doubtful. However, understanding the applicability of multi-headed architectures can help design and develop networks for deep learning that are more contextual and applicable to their problem domains.<\/p>\n
References<\/h2>\n\n- Caruana, R. Multitask Learning. Machine Learning 28, 41\u201375 (1997). https:\/\/doi.org\/10.1023\/A:1007379606734<\/li>\n
- Vaswani, A., et. al. Attention Is All You Need. arXiv.
\nhttps:\/\/doi.org\/10.48550\/arXiv.1706.03762<\/li>\n<\/ol>\n","protected":false},"excerpt":{"rendered":"Neural Networks have a diverse range of design architectures. These are often uniquely suited to specific problem domains or performance requirements. The Multi-Head design offers both semantic and computational isolation of elements of the network. This offers benefits both in model performance and development workflow. While this concept is not new to the world of […]<\/p>\n","protected":false},"author":2,"featured_media":3530,"comment_status":"closed","ping_status":"closed","sticky":false,"template":"","format":"standard","meta":{"_acf_changed":false,"inline_featured_image":false,"footnotes":""},"categories":[186,246],"tags":[745,742,743,741,744],"acf":[],"_links":{"self":[{"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/posts\/3527"}],"collection":[{"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/users\/2"}],"replies":[{"embeddable":true,"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/comments?post=3527"}],"version-history":[{"count":5,"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/posts\/3527\/revisions"}],"predecessor-version":[{"id":3533,"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/posts\/3527\/revisions\/3533"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/media\/3530"}],"wp:attachment":[{"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/media?parent=3527"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/categories?post=3527"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/www.alpharithms.com\/wp-json\/wp\/v2\/tags?post=3527"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}