* Lew et al. (2016) Lew, B. W. P., Apai, D., Zhou, Y., et al. 2016, ApJ, 829, L32, doi: 10.3847/2041-8205/829/2/L32
  * Lindegren et al. (2021) Lindegren, L., Klioner, S. A., Hernández, J., et al. 2021, A&A, 649, A2, doi: 10.1051/0004-6361/202039709
  * Liu et al. (2016) Liu, M. C., Dupuy, T. J., & Allers, K. N. 2016, ApJ, 833, 96, doi: 10.3847/1538-4357/833/1/96
  * Liu et al. (2010) Liu, M. C., Dupuy, T. J., & Leggett, S. K. 2010, ApJ, 722, 311, doi: 10.1088/0004-637X/722/1/311
  * Liu et al. (2006) Liu, M. C., Leggett, S. K., Golimowski, D. A., et al. 2006, ApJ, 647, 1393, doi: 10.1086/505561
  * Liu et al. (2013) Liu, M. C., Magnier, E. A., Deacon, N. R., et al. 2013, ApJ, 777, L20, doi: 10.1088/2041-8205/777/2/L20
  * Lodders & Fegley (2006) Lodders, K., & Fegley, Jr., B. 2006, Chemistry of Low Mass Substellar Objects (Springer Praxis Books), 1, doi: 10.1007/3-540-30313-8_1
  * Looper et al. (2007) Looper, D. L., Kirkpatrick, J. D., & Burgasser, A. J. 2007, AJ, 134, 1162, doi: 10.1086/520645
  * Looper et al. (2008) Looper, D. L., Kirkpatrick, J. D., Cutri, R. M., et al. 2008, ApJ, 686, 528, doi: 10.1086/591025
  * Luger et al. (2019) Luger, R., Agol, E., Foreman-Mackey, D., et al. 2019, AJ, 157, 64, doi: 10.3847/1538-3881/aae8e5
  * Luhman (2013) Luhman, K. L. 2013, ApJ, 767, L1, doi: 10.1088/2041-8205/767/1/L1
  * Luhman & Sheppard (2014) Luhman, K. L., & Sheppard, S. S. 2014, ApJ, 787, 126, doi: 10.1088/0004-637X/787/2/126
  * Luhman et al. (2007) Luhman, K. L., Patten, B. M., Marengo, M., et al. 2007, ApJ, 654, 570, doi: 10.1086/509073
  * Mace et al. (2013) Mace, G. N., Kirkpatrick, J. D., Cushing, M. C., et al. 2013, ApJS, 205, 6, doi: 10.1088/0067-0049/205/1/6
  * Madhusudhan et al. (2014) Madhusudhan, N., Knutson, H., Fortney, J. J., & Barman, T. 2014, in Protostars and Planets VI, ed. H. Beuther, R. S. Klessen, C. P. Dullemond, & T. Henning, 739, doi: 10.2458/azu_uapress_9780816531240-ch032
  * Manjavacas et al. (2013) Manjavacas, E., Goldman, B., Reffert, S., & Henning, T. 2013, A&A, 560, A52, doi: 10.1051/0004-6361/201321720
  * Manjavacas et al. (2016) Manjavacas, E., Goldman, B., Alcalá, J. M., et al. 2016, MNRAS, 455, 1341, doi: 10.1093/mnras/stv2048
  * Manjavacas et al. (2019) Manjavacas, E., Apai, D., Zhou, Y., et al. 2019, AJ, 157, 101, doi: 10.3847/1538-3881/aaf88f
  * Marley et al. (2012) Marley, M. S., Saumon, D., Cushing, M., et al. 2012, ApJ, 754, 135, doi: 10.1088/0004-637X/754/2/135
  * Metchev et al. (2008) Metchev, S. A., Kirkpatrick, J. D., Berriman, G. B., & Looper, D. 2008, ApJ, 676, 1281, doi: 10.1086/524721
  * Metchev et al. (2015) Metchev, S. A., Heinze, A., Apai, D., et al. 2015, ApJ, 799, 154, doi: 10.1088/0004-637X/799/2/154
  * Millar-Blanchaer et al. (2020) Millar-Blanchaer, M. A., Girard, J. H., Karalidi, T., et al. 2020, ApJ, 894, 42, doi: 10.3847/1538-4357/ab6ef2
  * Morales-Calderón et al. (2006) Morales-Calderón, M., Stauffer, J. R., Kirkpatrick, J. D., et al. 2006, ApJ, 653, 1454, doi: 10.1086/507866
  * Ochsenbein et al. (2000) Ochsenbein, F., Bauer, P., & Marcout, J. 2000, A&AS, 143, 23, doi: 10.1051/aas:2000169
  * Pineda et al. (2016) Pineda, J. S., Hallinan, G., Kirkpatrick, J. D., et al. 2016, ApJ, 826, 73, doi: 10.3847/0004-637X/826/1/73
  * Radigan (2014) Radigan, J. 2014, ApJ, 797, 120, doi: 10.1088/0004-637X/797/2/120
  * Radigan et al. (2012) Radigan, J., Jayawardhana, R., Lafrenière, D., et al. 2012, ApJ, 750, 105, doi: 10.1088/0004-637X/750/2/105
  * Radigan et al. (2014) Radigan, J., Lafrenière, D., Jayawardhana, R., & Artigau, E. 2014, ApJ, 793, 75, doi: 10.1088/0004-637X/793/2/75
  * Rayner et al. (2003) Rayner, J. T., Toomey, D. W., Onaka, P. M., et al. 2003, PASP, 115, 362, doi: 10.1086/367745
  * Reid et al. (2008) Reid, I. N., Cruz, K. L., Kirkpatrick, J. D., et al. 2008, AJ, 136, 1290, doi: 10.1088/0004-6256/136/3/1290
  * Reid et al. (2006a) Reid, I. N., Lewitus, E., Allen, P. R., Cruz, K. L., & Burgasser, A. J. 2006a, AJ, 132, 891, doi: 10.1086/505626
  * Reid et al. (2006b) Reid, I. N., Lewitus, E., Burgasser, A. J., & Cruz, K. L. 2006b, ApJ, 639, 1114, doi: 10.1086/499484
  * Reid et al. (1999) Reid, I. N., Kirkpatrick, J. D., Liebert, J., et al. 1999, ApJ, 521, 613, doi: 10.1086/307589
  * Robert et al. (2016) Robert, J., Gagné, J., Artigau, É., et al. 2016, ApJ, 830, 144, doi: 10.3847/0004-637X/830/2/144
  * Sahlmann et al. (2020) Sahlmann, J., Burgasser, A. J., Bardalez Gagliuffi, D. C., et al. 2020, MNRAS, 495, 1136, doi: 10.1093/mnras/staa1235
  * Schmidt et al. (2007) Schmidt, S. J., Cruz, K. L., Bongiorno, B. J., Liebert, J., & Reid, I. N. 2007, AJ, 133, 2258, doi: 10.1086/512158
  * Schneider et al. (2014) Schneider, A. C., Cushing, M. C., Kirkpatrick, J. D., et al. 2014, AJ, 147, 34, doi: 10.1088/0004-6256/147/2/34
  * Schneider et al. (2016a) Schneider, A. C., Greco, J., Cushing, M. C., et al. 2016a, ApJ, 817, 112, doi: 10.3847/0004-637X/817/2/112
  * Schneider et al. (2018) Schneider, A. C., Hardegree-Ullman, K. K., Cushing, M. C., Kirkpatrick, J. D., & Shkolnik, E. L. 2018, AJ, 155, 238, doi: 10.3847/1538-3881/aabfc2
  * Schneider et al. (2017) Schneider, A. C., Windsor, J., Cushing, M. C., Kirkpatrick, J. D., & Shkolnik, E. L. 2017, AJ, 153, 196, doi: 10.3847/1538-3881/aa6624
  * Schneider et al. (2016b) Schneider, A. C., Windsor, J., Cushing, M. C., Kirkpatrick, J. D., & Wright, E. L. 2016b, ApJ, 822, L1, doi: 10.3847/2041-8205/822/1/L1
  * Scholz (2010) Scholz, R.-D. 2010, A&A, 510, L8, doi: 10.1051/0004-6361/201014078
  * Sheppard & Cushing (2009) Sheppard, S. S., & Cushing, M. C. 2009, AJ, 137, 304, doi: 10.1088/0004-6256/137/1/304
  * Showman et al. (2019) Showman, A. P., Tan, X., & Zhang, X. 2019, ApJ, 883, 4, doi: 10.3847/1538-4357/ab384a
  * Song et al. (2012) Song, I., Zuckerman, B., & Bessell, M. S. 2012, AJ, 144, 8, doi: 10.1088/0004-6256/144/1/8
  * Tan & Showman (2021a) Tan, X., & Showman, A. P. 2021a, MNRAS, 502, 2198, doi: 10.1093/mnras/stab097
  * Tan & Showman (2021b) —. 2021b, MNRAS, 502, 678, doi: 10.1093/mnras/stab060
  * Thompson et al. (2013) Thompson, M. A., Kirkpatrick, J. D., Mace, G. N., et al. 2013, PASP, 125, 809, doi: 10.1086/671426
  * Tinney et al. (2003) Tinney, C. G., Burgasser, A. J., & Kirkpatrick, J. D. 2003, AJ, 126, 975, doi: 10.1086/376481
  * Tremblin et al. (2016) Tremblin, P., Amundsen, D. S., Chabrier, G., et al. 2016, ApJ, 817, L19, doi: 10.3847/2041-8205/817/2/L19
  * Tremblin et al. (2015) Tremblin, P., Amundsen, D. S., Mourier, P., et al. 2015, ApJ, 804, L17, doi: 10.1088/2041-8205/804/1/L17
  * Tremblin et al. (2020) Tremblin, P., Phillips, M. W., Emery, A., et al. 2020, A&A, 643, A23, doi: 10.1051/0004-6361/202038771
  * Tremblin et al. (2017) Tremblin, P., Chabrier, G., Baraffe, I., et al. 2017, ApJ, 850, 46, doi: 10.3847/1538-4357/aa9214
  * Tremblin et al. (2019) Tremblin, P., Padioleau, T., Phillips, M. W., et al. 2019, ApJ, 876, 144, doi: 10.3847/1538-4357/ab05db
  * Vacca et al. (2003) Vacca, W. D., Cushing, M. C., & Rayner, J. T. 2003, PASP, 115, 389, doi: 10.1086/346193
  * Vos et al. (2019a) Vos, J., Allers, K., Apai, D., et al. 2019a, BAAS, 51, 253
  * Vos et al. (2017) Vos, J. M., Allers, K. N., & Biller, B. A. 2017, ApJ, 842, 78, doi: 10.3847/1538-4357/aa73cf
  * Vos et al. (2018) Vos, J. M., Allers, K. N., Biller, B. A., et al. 2018, MNRAS, 474, 1041, doi: 10.1093/mnras/stx2752
  * Vos et al. (2022) Vos, J. M., Faherty, J. K., Gagné, J., et al. 2022, ApJ, 924, 68, doi: 10.3847/1538-4357/ac4502
  * Vos et al. (2019b) Vos, J. M., Biller, B. A., Bonavita, M., et al. 2019b, MNRAS, 483, 480, doi: 10.1093/mnras/sty3123
  * Vos et al. (2020) Vos, J. M., Biller, B. A., Allers, K. N., et al. 2020, AJ, 160, 38, doi: 10.3847/1538-3881/ab9642
  * Vrba et al. (2004) Vrba, F. J., Henden, A. A., Luginbuhl, C. B., et al. 2004, AJ, 127, 2948, doi: 10.1086/383554
  * Wenger et al. (2000) Wenger, M., Ochsenbein, F., Egret, D., et al. 2000, A&AS, 143, 9, doi: 10.1051/aas:2000332
  * Wright et al. (2010) Wright, E. L., Eisenhardt, P. R. M., Mainzer, A. K., et al. 2010, AJ, 140, 1868, doi: 10.1088/0004-6256/140/6/1868
  * Zhang et al. (2021) Zhang, Z., Liu, M. C., Best, W. M. J., Dupuy, T. J., & Siverd, R. J. 2021, ApJ, 911, 7, doi: 10.3847/1538-4357/abe3fa
  * Zuckerman (2019) Zuckerman, B. 2019, ApJ, 870, 27, doi: 10.3847/1538-4357/aaee66


# Meta-Reinforcement Learning for Adaptive Control of Second Order Systems

Daniel G. McClement1, Nathan P. Lawrence,2, Michael G. Forbes3,
Philip D. Loewen2, Johan U. Backström4, R. Bhushan Gopaluni1 1Department of
Chemical and Biological Engineering, University of British Columbia,
Vancouver, BC Canada2Department of Mathematics, University of British
Columbia, Vancouver BC, Canada3Honeywell Process Solutions, North Vancouver,
BC Canada4Backstrom Systems Engineering Ltd., North Vancouver, BC Canada©2022
the authors. This work has been accepted to AdCONIP 2022. We gratefully
acknowledge the financial support from Natural Sciences and Engineering
Research Council of Canada (NSERC) and Honeywell Connected Plant.

###### Abstract

Meta-learning is a branch of machine learning which aims to synthesize data
from a distribution of related tasks to efficiently solve new ones. In process
control, many systems have similar and well-understood dynamics, which
suggests it is feasible to create a generalizable controller through meta-
learning. In this work, we formulate a meta reinforcement learning (meta-RL)
control strategy that takes advantage of known, offline information for
training, such as a model structure. The meta-RL agent is trained over a
distribution of model parameters, rather than a single model, enabling the
agent to automatically adapt to changes in the process dynamics while
maintaining performance. A key design element is the ability to leverage
model-based information offline during training, while maintaining a model-
free policy structure for interacting with new environments. Our previous work
has demonstrated how this approach can be applied to the industrially-relevant
problem of tuning proportional-integral controllers to control first order
processes. In this work, we briefly reintroduce our methodology and
demonstrate how it can be extended to proportional-integral-derivative
controllers and second order systems.

##  I Introduction

_Meta-learning_ , or “learning to learn”, is an active area of research in
which the objective is to learn an underlying structure governing a
distribution of possible tasks [2]. This framework can be applied to
reinforcement learning (RL) for learning a generalized goal-oriented “policy”
over a large number of stochastic environments [3]. In process control
applications, meta reinforcement learning (meta-RL) is appealing because many
systems have similar dynamics or a known structure [4], which suggests
training a policy over a distribution of systems could enable more rapid
online deployment without specifying a system-specific model. Moreover,
extensive online learning is impractical for training over a large number of
systems; by focusing on learning a underlying structure across systems with
different process dynamics, we can more readily adapt to a new system.

This paper proposes a method for improving the _online_ sample efficiency of
RL agents. Our approach is to train a meta-RL agent _offline_ by exposing it
to a broad distribution of different dynamics (or “tasks”). The agent
synthesizes its experience from different environments to quickly learn an
optimal policy for its present environment. The basis for our strategy is a
recurrent neural network (RNN) policy. The hidden state of the RNN serves as
an encoding of the system dynamics, which provides the network with “context”
for its policy. We aim to use this framework to develop a universal controller
which can quickly adapt to effectively control any process rather than a
single task. The training is performed completely offline and the result is a
single RL agent that can quickly adapt its policy to a new environment in a
model-free fashion.

We apply this general method to the industrially-relevant problem of
autonomous controller tuning and adaptation. We show how our trained agent can
adaptively fine-tune proportional-integral-derivative (PID) controller
parameters when the underlying dynamics drift or are not precisely known.
Perhaps the most appealing consequence of this method is that it removes the
need to accommodate a training algorithm on a system-by-system basis — for
example, through extensive online training or transfer learning,
hyperparameter tuning, or system identification — because the adaptive policy
is pre-computed and represented in a single model.

This paper extends our previous work which focused only on systems which could
be modeled as first order plus time delay and PI controllers [1]. In
particular, in Section III we extend our framework to second order systems and
PID controllers. Section IV provides new simulation studies and analysis of
the trained meta-RL agent.

##  II Background

###  II-A Reinforcement learning

In this section, we give a brief overview of deep RL. We refer the reader to
[5, 6], for tutorial overviews of deep RL with applications to process
control. We use the standard RL terminology that can be found in [7]. [2]
gives a unified survey of deep meta-learning.

The RL framework consists of an _agent_ and an _environment_. For each _state_
$s_{t}\in\mathcal{S}$ (the state-space) the agent encounters, it takes some
_action_ $a_{t}\in\mathcal{A}$ (the action-space), leading to a new state
$s_{t+1}$. The action is chosen according to a conditional probability
distribution $\pi$ called a _policy_ ; we denote this relationship by
$a_{t}\sim\pi(a_{t}|s_{t})$. Although the system dynamics are not necessarily
known, we assume they can be described as a Markov decision process (MDP) with
initial distribution $p(s_{0})$ and transition probability
$p(s_{t+1}|s_{t},a_{t})$. A state-space model in control is a special case of
an MDP, where the states are the (minimal realization) vector that
characterizes the system, while the actions are the control inputs. At each
time step, a bounded scalar _cost_ $c_{t}=c(s_{t},a_{t})$ is evaluated. The
cost function describes the desirability of a state-action pair and defining
it is a key part of the design process. The overall objective, however, is the
expected long-term cost. In terms of a user-specified discount factor
$0<\gamma<1$, the optimization problem of interest becomes

minimize $\displaystyle J(\theta)=\mathbb{E}_{h\sim
p^{\pi_{\theta}}(\cdot)}\left[\sum_{t=1}^{\infty}\gamma^{t-1}c(s_{t},\pi_{\theta}(s_{t}))\middle|s_{0}\right]$
(1) over all $\displaystyle\theta\in\mathbb{R}^{n}.$

In this problem, $h\sim p^{\pi}$ refers to a typical trajectory
$h~{}=~{}(s_{0},a_{0},c_{0},,\ldots,s_{N},a_{N},c_{N})$ generated by the
policy $\pi$ with subsequent states distributed according to $p$. Within the
space of all possible policies, we optimize over a parameterized subset whose
members are denoted $\pi_{\theta}$. We use $\theta$ as a generic vector of
parameters.

Common approaches to solving Problem (1) involve techniques based on
$Q$-learning (value-based methods) and the policy gradient theorem (policy-
based methods) [7], or a combination of both called _actor-critic_ methods
[8]. Closely-related functions to $J$ are the $Q$-function (state-action value
function) and value function, respectively:

$\displaystyle Q(s_{t},a_{t})$ $\displaystyle=\mathbb{E}_{h\sim
p^{\pi}(\cdot)}\left[\sum_{k=t}^{\infty}\gamma^{t-1}c(s_{k},a_{k})\middle|s_{t},a_{t}\right]$
(2) $\displaystyle V(s_{t})$ $\displaystyle=\mathbb{E}_{h\sim
p^{\pi}(\cdot)}\left[\sum_{k=t}^{\infty}\gamma^{t-1}c(s_{k},a_{k})\middle|s_{t}\right].$
(3)

The _advantage function_ is then $A(s,a)=Q(s,a)-V(s)$. These functions help
form the basis for deep RL algorithms, that is, algorithms that use deep
neural networks to solve RL tasks. Deep neural networks are a flexible form of
function approximators, well-suited for learning complex control laws.
Moreover, function approximation methods make RL problems tractable in
continuous state and action spaces [9, 10].

A standard approach to solving Problem (1) uses gradient descent:

$\theta\leftarrow\theta-\alpha\nabla J(\theta),$ (4)

where $\alpha>0$ is a step-size parameter.

Analytic expressions for such a gradient exist for both stochastic and
deterministic policies [7, 9]. However, in practice, approximations are
necessary. We use the proximal policy optimization (PPO) reinforcement
learning algorithm to optimize Problem (1) using a surrogate objective to
compute (4) based on the advantage function [11]. We use PPO because of its
sample efficiency and robustness to the choice of model hyperparameters. Full
implementation and technical details can be found in [1, 11].

###  II-B Meta reinforcement learning

Figure 1: A diagram of the meta-RL agent’s interactions with the task
distribution. Each MDP represents a different task sampled from
$p(\mathcal{T})$.

While the RL framework mentioned above can achieve impressive results in a
wide range of domains, it is formulated for a single MDP. In contrast, meta-RL
aims to generalize agents to a distribution of MDPs. Formally, a single MDP
can be characterized by a tuple
$\mathcal{T}=(\mathcal{S},\mathcal{A},p,c,\gamma)$; in contrast, meta-RL
tackles an optimization problem over a distribution
$p_{\text{meta}}(\mathcal{T})$ of MDPs. Therefore, in the meta-RL terminology,
a “task” is simply all the components comprising a single RL problem. The
problem of interest in the meta-RL setting is a generalization of the standard
RL objective in Problem (1) [2]:

minimize $\displaystyle
J_{\text{meta}}(\bm{\Theta})=\mathbb{E}_{\mathcal{T}\sim
p_{\text{meta}}(\mathcal{T})}\left[J(\theta^{*}(\mathcal{T},\bm{\Theta}))\right]$
(5) over all $\displaystyle\bm{\Theta}\in\mathbb{R}^{n}.$

Crucially, in the context of process control, meta-RL does _not_ aim to find a
single controller that performs well across different plants. Note that
$\theta^{*}$ in Eq. 5 is the optimal weight vector in Eq. 1 as a function of a
sampled MDP $\mathcal{T}$ and the meta-weights $\bm{\Theta}$. Meta-RL agents
aim to simultaneously learn the underlying structure characterizing different
plants and the corresponding optimal control strategy under its cost function.
The practical benefit is that this enables RL agents to quickly adapt to novel
environments.

There are two components to meta-learning algorithms: the model (e.g., the
policy network) that solves a given task, and a set of meta-parameters that
learn how to update the model [12, 13]. Due to the shared structure among
tasks in process control applications, we are interested in _context-based_
meta-RL methods [14, 15, 16]. These approaches learn a latent representation
of each task, enabling the agent to simultaneously learn the context and the
policy for a given task. This contrasts with other recent RL work on PID
tuning, such as [17], where a process model is identified on a case-by-case
basis and an RL agent is trained for a specific process.

Our method is similar to [14]: the main idea is to view Eq. 5 as a single RL
problem. For each MDP $\mathcal{T}\sim p(\mathcal{T})$, the meta-RL agent has
a maximum number of time steps, $T$, to interact with the environment, called
an _episode_. As each episode progresses, the RL agent has an internal hidden
state vector $z_{t}$ which evolves with each time step through the MDP based
on the RL states the agent observes: $z_{t}=f_{\bm{\Theta}}(z_{t-1},s_{t})$.
The RL agent conditions its actions on both $s_{t}$ and $z_{t}$. An
illustration of this concept is shown in Fig. 1. Therefore, the purpose of the
meta-parameters $\bm{\Theta}$ is to quickly adapt a control policy for a MDP
$\mathcal{T}\sim p(\mathcal{T})$ by solving for a suitable set of MDP-specific
parameters encoded by $z_{t}$. This is why this approach is described as meta-
RL; rather than training a reinforcement learning agent to control a process,
we are training a meta-reinforcement learning agent to find a suitable set of
parameters for a reinforcement learning agent which can control a process. The
advantage of training a meta-RL agent is that the final model is capable of
controlling every MDP across the task distribution $p(\mathcal{T})$ whereas a
regular RL agent could only be optimized for a single task $\mathcal{T}$.

Clearly, the key component of the above framework is the hidden state. This is
generated with a recurrent neural network (RNN), a special neural network
structure for processing sequential data. An RNN can be thought of as a
nonlinear state-space system that is optimized for some objective. The RNN
structure we use in practice is a gated recurrent network [18].

##  III Meta-RL for process control

We apply the meta-RL framework to the problem of tuning proportional-integral-
derivative (PID) controllers. The formulation can be applied to any fixed-
structure controller, but we focus on PID controllers due to their industrial
prevalence.

###  III-A Tasks, states, actions, costs

The systems of interest are second-order plus time delay (SOPTD):

$G(s)=\frac{K}{(\tau_{1}s+1)(\tau_{2}s+1)}e^{-\theta s},$ (6)

where $K$ is the process gain, $\tau_{1}$ and $\tau_{2}$ are the system poles
($\tau_{1}\geq\tau_{2})$, $\theta$ is the time delay, and $s$ is the Laplace
variable (not to be confused with $s_{t}$, which represents the RL state at
time step $t$).

A PID controller has the form:

$C(s)=K_{c}\left(1+\frac{1}{\tau_{I}s}+\tau_{D}s\right),$ (7)

where $K_{c}$, $\tau_{I}$, $\tau_{D}$ are tuning parameters. We also apply a
low-pass filter to the derivative error term in the controller. Prior work on
RL for PID tuning suggests an update scheme of the form [19]:

$\displaystyle[K_{c},\tau_{I},\tau_{D}]$
$\displaystyle\leftarrow[K_{c},\tau_{I},\tau_{D}]+\alpha\nabla
J([K_{c},\tau_{I},\tau_{D}])$ (8)
$\displaystyle=[K_{c},\tau_{I},\tau_{D}]+\Delta[K_{c},\tau_{I},\tau_{D}]$ (9)

where the RL policy is directly parameterized as a PID controller. We follow
this update scheme and take the actions to be changes to the PID parameters
$\Delta[K_{c},\tau_{I},\tau_{D}]$.

The MDP state contains the current PID parameters as well as the proportional
setpoint error and the integral setpoint error calculated from the beginning
of an episode, $t_{0}$, to the current time step, $t$.

$s_{t}=\left[K_{c},\tau_{I},\tau_{D},e_{t},\int_{t_{0}}^{t}e_{\tau}d\tau\right]$
(10)

The RL agent is trained to minimize its discounted future cost interacting
with different tasks. The cost function used to train the meta-RL agent is the
squared error from a target trajectory, shown in Eq. 11. The target trajectory
is calculated by applying a first order filter to the set point signal, shown
in Eq. 12. The time constant of this filter is set to the desired closed-loop
time constant, $\tau_{cl}$. A target closed-loop time constant is chosen for
robustness and smooth control action following the PID tuning rules presented
in [20] and shown in Eq. 13. An $L1$ regularization penalty $\beta>0$ on the
agent’s actions is also added to the cost function to encourage sparsity in
the meta-RL agent’s output and help the tuning algorithm converge to a
constant set of PID parameters (rather than acting as a non-linear feedback
controller and constantly changing the controller parameters in response to
the current low-level state of the system).

$c_{t}=(y_{desired,t}-y_{t})^{2}+\beta_{1}|\Delta
K_{c}|+\beta_{2}|\Delta\tau_{I}|+\beta_{3}|\Delta\tau_{D}|,$ (11)

where

$\displaystyle Y_{desired}(s)$
$\displaystyle=\frac{y_{sp}}{\tau_{cl}s+1}e^{-\theta s}$ (12)
$\displaystyle\tau_{cl}$ $\displaystyle=2\tau_{1}+\tau_{2},~{}\text{with
}\tau_{1}>\tau_{2}$ (13)

Comparing the RL state definition to the RL cost definition, we see similar
trajectories through different MDPs will receive very different costs
depending on the underlying system dynamics in the particular tasks being
controlled. In order for the meta-RL agent to perform well on a new task, it
needs to perform implicit system identification to generate an internal
representation of the system dynamics.

The advantages of this meta-RL scheme for PID tuning are summarized as
follows:

  * •

Tuning is performed in closed-loop and without explicit system identification.

  * •

Tuning is performed automatically even as the underlying system changes.

  * •

The agent can be deployed on novel “in distribution” systems (e.g. systems
within the task distribution $p(\mathcal{T})$) without any online training.
Additionally, nearly any system’s data can be modified to be “in-distribution”
[1].

  * •

The meta-RL agent is a single model that is trained once, offline, meaning one
does not need to specify hyperparameters on a task-by-task basis.

  * •

The meta-RL agent’s cost function is conditioned on the process dynamics and
will produce consistent closed-loop control behaviour on different systems.

This approach is not limited to PID tuning. It can also be applied to other
scenarios where the model _structure_ is known. The agent then learns to
behave near-optimally inside each task in the training distribution, bypassing
the need to identify model parameters and only train on that instance of the
dynamics.

###  III-B RL agent structure

The structure of the meta-RL agent is shown in Fig. 2. The grey box shows the
“actor”, i.e., the part of the agent used online for controller tuning.
Through interacting with a system and observing the RL states at each time
step, the agent’s recurrent layers create an embedding (hidden state) which
encodes information needed to tune the PID parameters, including information
about the system dynamics and the uncertainty associated with this
information. These embeddings essentially represent process-specific RL
parameters which are updated as the meta-RL agent’s knowledge of the process
dynamics changes. Two fully connected layers use these embeddings to recommend
adjustments to the controller’s PID parameters. The inclusion of recurrent
layers is essential for the meta-RL agent’s performance. Having a hidden state
carried between time steps equips the agent with memory and enables the agent
to learn a representation of the process dynamics. A traditional feedforward
RL network would be unable to differentiate between different tasks and would
perform significantly worse. This concept is demonstrated in [21].

Outside of the grey box are additional parts of the meta-RL agent which are
only used during offline training. The “critic” (shown in green) is trained to
calculate the value (an estimate of the agent’s discounted future cost in the
current MDP given the current RL state). This value function is used to train
the meta-RL actor through gradient descent via the PPO algorithm.

A unique strategy we use to improve the training efficiency of the meta-RL
agent is to give the critic network access to “privileged information”,
defined as any additional information outside the RL state and denoted as
$\zeta$. In addition to the RL state, the critic conditions its estimates of
the value function on the true process parameters ($K$, $\tau_{1}$,
$\tau_{2}$, and $\theta$), as well as the deep hidden state — the hidden state
of the second recurrent layer — of the actor. Knowledge of a task’s process
dynamics, as well as knowledge of the actor’s internal representation of the
process dynamics through its hidden state, allows the critic to more
accurately estimate the value function, which improves the quality of the
surrogate objective function used to train the actor and the training
efficiency of the algorithm [1]. Equipping the critic with this information
also allows it to operate as a simpler feedforward neural network rather than
a recurrent network like the actor.

The privileged information given to the critic network may at first appear to
conflict with the advantages of the proposed meta-RL tuning method, since the
critic requires the true system parameters and much simpler tuning methods for
PID controllers exist if such information is known. However, this information
is only required during _offline_ training. The meta-RL agent is trained on
simulated systems with known process dynamics, but the end result of this
training procedure is a meta-RL agent that can be used to tune PID parameters
for a real process _online_ with no task-specific training or knowledge of the
process dynamics. The portion of the meta-RL agent operating online contained
in the grey box only requires RL state information — process data — at each
time step.

Figure 2: The structure of the RL agent. The control policy used online is
shown in the grey box while the critic used during offline training is shown
in green.

###  III-C Training procedure

The meta-RL agent is trained by uniformly sampling $K$, $\tau_{1}$,
$\tau_{2}$, and $\theta$ to create a SOPTD system and initializing a PID
controller with $K_{c}=0.05$, $\tau_{I}=1.0$, $\tau_{D}=0.2$. The state of the
system is randomly initialized near zero by sampling from $\mathcal{N}(0,0.1)$
and the set point is switched between $1$ and $-1$ every $11$ units of time.
The meta-RL agent has no inherent time scale and so we keep the units of time
general to highlight the applicability of the proposed PID tuning algorithm to
both fast and slow processes (allowing time constants on the order of
milliseconds or hours).

We consider the following sets of SOPTD model parameters in our experiments:
$K$ and $\tau_{1}$ are sampled from the interval $[0.25,1.0]$; The ratios
$\frac{\tau_{2}}{\tau_{1}}$ and $\frac{\theta}{\tau_{1}}$ are sampled from the
interval $[0.0,1.0]$. In [1] we demonstrate how training across this range of
parameters can be quite versatile in practice using data augmentation to
adjust the dynamics of most second order processes into this distribution from
the meta-RL agent’s perspective. The PPO algorithm is adapted from Open AI’s
“Spinning Up” implementation and modified to accommodate the inclusion of a
recurrent neural network and distribution of control tasks [22].

##  IV Results

After training, the meta-RL model is tested on a range of processes within
$p(\mathcal{T})$. To visualize the performance of the meta-RL agent, we plot
the mean squared error from the target trajectory when the set point changes 2
units for different SOPTD systems after the agent has had time to adapt.
Because the set of possible processes is four-dimensional, we hold two
parameters constant at their mean values to plot how the performance varies as
the other two parameters are manipulated. The performance is shown in Fig. 3.

Figure 3: Performance of the Meta-RL agent across $p(\mathcal{T})$. By
default, $K=0.5$, $\tau_{1}=0.5$, $\dfrac{\tau_{2}}{\tau_{1}}=0.5$,
$\dfrac{\theta}{\tau_{1}}=0.5$. The color represents the mean squared error
from a target trajectory during a step response, as shown in Fig. 4.

In Fig. 3 we see the agent replicates the desired trajectory quite accurately
and consistently across $p(\mathcal{T})$. Performance is significantly worse
when the process gain and dominant pole both have small magnitudes. These
systems require the largest changes to the initial PID parameters and the
regularization added in Eq. 11 to changes in PID parameters results in a
slightly under-tuned controller in these cases.

To further illustrate the meta-RL agent’s PID tuning performance, a set point
change on the best and worst-case systems from Fig. 3 is shown in Fig. 4. A
third trajectory is shown to illustrate the performance on systems which have
more prominent second-order behaviour ($\tau_{2}\approx\tau_{1})$.

Figure 4: A demonstration of the trajectory-tracking of PID tunings produced
by the meta-RL agent ability to track the target trajectory in the worst case
(left), best case (center) and with significant second order dynamics (right).

Fig. 5 demonstrates the meta-RL agent’s initial PID adaptation when
initialized on a new system. To illustrate the meta-RL agent’s ability to
adapt its tunings when the system dynamics change, more prominent second-order
behaviour is introduced by increasing $\tau_{2}$ and we see the meta-RL agent
is able to recognize and adapt the PID parameters in response to this change
in the process dynamics.

Figure 5: A demonstration of the initial PID tunings produced by the meta-RL
agent on the system $k=0.5,~{}\tau_{1}=0.8,~{}\tau_{2}=0.1,~{}\theta=0.05$ as
well as the meta-RL agent’s ability to adapt the PID parameters when the
process dynamics change ($\tau_{2}$ is increased to 0.8 to simulate more
prominent second-order behaviour).

Next, we analyze the hidden state of the meta-RL agent after adapting to a
process. We reduce the hidden state to two principal components (PCs) for two-
dimensional visualization in Fig. 6. We see the hidden state of the meta-RL
agent does encode information about the process gain and dominant pole
$\tau_{1}$, confirming our hypothesis that the agent is performing implicit
system identification.

Figure 6: Principal component visualization of the meta-RL agent’s deep hidden
state after adapting to various processes.

In Fig. 7, we recalculate the principal components when $K$ and $\tau_{1}$ are
held constant to better visualize whether the hidden state contains
information about $\tau_{2}$ and $\theta$ as well. From the pattern, we see
information about $\tau_{2}$ and $\theta$ are encoded in the hidden states.
Another interesting result is the differentiation between processes with first
or near first-order dynamics ($\tau_{2}=0$ or $\tau_{2}\ll\theta$) from other
processes within $p(\mathcal{T})$.

Figure 7: Principal component visualization of the meta-RL agent’s deep hidden
state after adapting to various processes with $K=1$ and $\tau_{1}=1$.

##  V Conclusion

In this work, we presented a meta-RL model capable of tuning fixed-structure
controllers in closed-loop without explicit system identification and
demonstrated our approach using PID controllers. The tuning algorithm can be
used to help automate the initial tuning of controllers or maintenance of
controllers by adaptively updating the controller parameters as process
dynamics change over time. Assuming the magnitude of the process gain and time
constant are known, the meta-RL tuning algorithm can be applied to any system
which can be reasonably approximated as a SOPTD system. The present work
focuses on lag dominant SOPTD systems ($\tau_{1}>\theta$), however the results
could be extended to dead-time dominant systems by expanding the task
distribution $p(\mathcal{T})$ used during training.

A major challenge of applying RL to industrial process control is sample
efficiency. The meta-RL model presented in this work addresses this problem by
training a model to control a large distribution of possible systems offline
in advance. The meta-RL model is then able to tune fixed-structure process
controllers online with no process-specific training and no process model.
There are two key design considerations which enable this performance. First
is the inclusion of a hidden state in the RL agent, giving the meta-RL agent a
memory it uses to learn internal representations of the process dynamics
through process data. Second is constructing a value function which uses extra
information in addition to the RL state.

## References

  * [1] D. G. McClement, N. P. Lawrence, J. U. Backström, P. D. Loewen, M. G. Forbes, and R. B. Gopaluni, “Meta-reinforcement learning for the tuning of PI controllers: An offline approach,” _Journal of Process Control_ , vol. 118, pp. 139–152, 2022.
  * [2] M. Huisman, J. N. van Rijn, and A. Plaat, “A survey of deep meta-learning,” _Artificial Intelligence Review_ , pp. 1–59, 2021.
  * [3] C. Finn, P. Abbeel, and S. Levine, “Model-agnostic meta-learning for fast adaptation of deep networks,” _arXiv preprint arXiv:1703.03400_ , 2017.
  * [4] C. Madhuranthakam, A. Elkamel, and H. Budman, “Optimal tuning of PID controllers for FOPTD, SOPTD and SOPTD with lead processes,” _Chemical Engineering and Processing: Process Intensification_ , vol. 47, no. 2, pp. 251–264, 2008.
  * [5] R. Nian, J. Liu, and B. Huang, “A review on reinforcement learning: Introduction and applications in industrial process control,” _Computers & Chemical Engineering_, p. 106886, 2020.
  * [6] S. Spielberg, A. Tulsyan, N. P. Lawrence, P. D. Loewen, and R. B. Gopaluni, “Toward self-driving processes: A deep reinforcement learning approach to control,” _AIChE Journal_ , 2019.
  * [7] R. S. Sutton and A. G. Barto, _Reinforcement learning: An introduction_. MIT press, 2018.
  * [8] V. R. Konda and J. N. Tsitsiklis, “Actor-critic algorithms,” in _Proceedings of the Advances in Neural Information Processing Systems_ , Denver, USA, 2000, pp. 1008–1014.
  * [9] D. Silver, G. Lever, N. Heess, T. Degris, D. Wierstra, and M. Riedmiller, “Deterministic policy gradient algorithms,” in _International conference on machine learning_. PMLR, 2014, pp. 387–395.
  * [10] R. S. Sutton, D. A. McAllester, S. P. Singh, and Y. Mansour, “Policy gradient methods for reinforcement learning with function approximation,” in _Proceedings of the Advances in Neural Information Processing Systems_ , 2000, pp. 1057–1063.
  * [11] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, “Proximal policy optimization algorithms,” _arXiv preprint arXiv:1707.06347_ , 2017\.
  * [12] S. Bengio, Y. Bengio, J. Cloutier, and J. Gescei, “On the optimization of a synaptic learning rule,” in _Optimality in Biological and Artificial Networks?_ Routledge, 2013, pp. 281–303.
  * [13] M. Andrychowicz, M. Denil, S. Gomez, M. W. Hoffman, D. Pfau, T. Schaul, B. Shillingford, and N. De Freitas, “Learning to learn by gradient descent by gradient descent,” in _Advances in neural information processing systems_ , 2016, pp. 3981–3989.
  * [14] Y. Duan, J. Schulman, X. Chen, P. L. Bartlett, I. Sutskever, and P. Abbeel, “$\text{RL}^{2}$: Fast reinforcement learning via slow reinforcement learning,” _arXiv preprint arXiv:1611.02779_ , 2016.
  * [15] J. X. Wang, Z. Kurth-Nelson, D. Tirumala, H. Soyer, J. Z. Leibo, R. Munos, C. Blundell, D. Kumaran, and M. Botvinick, “Learning to reinforcement learn,” _arXiv preprint arXiv:1611.05763_ , 2016.
  * [16] K. Rakelly, A. Zhou, D. Quillen, C. Finn, and S. Levine, “Efficient off-policy meta-reinforcement learning via probabilistic context variables,” in _International conference on machine learning_ , 2019, pp. 5331–5340.
  * [17] O. Dogru, K. Velswamy, F. Ibrahim, Y. Wu, A. S. Sundaramoorthy, B. Huang, S. Xu, M. Nixon, and N. Bell, “Reinforcement learning approach to autonomous PID tuning,” _Computers & Chemical Engineering_, vol. 161, p. 107760, 2022.
  * [18] K. Cho, B. Van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio, “Learning phrase representations using rnn encoder-decoder for statistical machine translation,” _arXiv preprint arXiv:1406.1078_ , 2014.
  * [19] N. P. Lawrence, M. G. Forbes, P. D. Loewen, D. G. McClement, J. U. Backström, and R. B. Gopaluni, “Deep reinforcement learning with shallow controllers: An experimental application to PID tuning,” _Control Engineering Practice_ , vol. 121, p. 105046, 2022.
  * [20] S. Skogestad, “Simple analytic rules for model reduction and PID controller tuning,” _Journal of process control_ , vol. 13, no. 4, pp. 291–309, 2003\.
  * [21] D. G. McClement, N. P. Lawrence, P. D. Loewen, M. G. Forbes, J. U. Backström, and R. B. Gopaluni, “A meta-reinforcement learning approach to process control,” _IFAC-PapersOnLine_ , vol. 54, no. 3, pp. 685–692, 2021\.
  * [22] J. Achiam, “Spinning Up in Deep Reinforcement Learning,” 2018.


# IGOT: Information Gain Optimized Tokenizer on Domain Adaptive Pretraining

Dawei Feng  Tsinghua University
Beijing, China
<EMAIL_ADDRESS>Yihai Zhang  Tsinghua University
Beijing, China
<EMAIL_ADDRESS>Zhixuan Xu  Tsinghua University
Beijing, China
<EMAIL_ADDRESS>
###### Abstract

Pretrained Large Language Models (LLM) such as ChatGPT, Claude, etc. have
demonstrated strong capabilities in various fields of natural language
generation. However, there are still many problems when using LLM in
specialized domain-specific fields. When using generative AI to process
downstream tasks, a common approach is to add new knowledge (e.g., private
domain knowledge, cutting-edge information) to a pretrained model through
continued training or fine-tuning. However, whether there is a universal
paradigm for domain adaptation training is still an open question. In this
article, we proposed Information Gain Optimized Tokenizer (IGOT), which
analyzes the special token set of downstream tasks, constructs a new subset
using heuristic function $\phi$ with the special token and its information
gain, to build new domain-specific tokenizer, and continues pretraining on the
downstream task data. We explored the many positive effects of this method’s
customized tokenizer on domain-adaptive pretraining and verified this method
can perform better than the ordinary method of just collecting data and fine-
tuning. Based on our experiment, the continued pretraining process of IGOT
with LLaMA-7B achieved 11.9% token saving, 12.2% training time saving, and
5.8% maximum GPU VRAM usage saving, combined with the T5 model, we can even
reach a 31.5% of training time saving, making porting general generative AI to
specific domains more effective than before. In domain-specific tasks,
supervised $IGOT_{\tau}$ shows great performance on reducing both the
convergence radius and convergence point during keep pretraining.

###### Index Terms:

Large Language Model, Customized Tokenizer, Domain Adaption Pretraining,
Information Entropy

##  I Introduction

Figure 1: An example of IGOT compared to LLaMA2 Tokenizer. Given the input
“Introduce Openlane, an EDA tool”, the original LLaMA2 tokenizer will split
into 13 tokens including meaningless parts such as “rodu” or “ane”, as shown
in Fig. 1(a). However, our IGOT method allows the tokenizer to process it into
8 meaningful parts, keep proprietary words like OpenLane, and reach a 38.46%
of space reduction.

After extensive training on publicly available web data, large language models
have exhibited remarkable capabilities in a wide range of common sense
reasoning tasks[1]. Some studies have indicated that these models, during
pretraining, not only learn contextual text vector representations and
probability distributions but are also likely to learn grammar[2], syntax[3],
factual knowledge[4], even behavior patterns strikingly similar to human
common sense[5].

While general-purpose LLMs trained on massive internet datasets demonstrate
impressive capabilities across diverse domains[6], recent work with domain-
specific models like BloombergGPT [7] and BioMedLM[8] highlights their
superior performance within specific fields.

Early efforts to build large-scale language models utilized n-grams and simple
smoothing techniques[9][10]. Various neural network architectures were later
applied to the language modeling task, including feedforward networks[11] and
recurrent networks[12]. The introduction of the Transformer architecture[13]
led to the development of highly scalable language models[14], demonstrating a
predictable relationship between language modeling loss and scaling factors
such as model size, number of training tokens, and compute budget.

The AI community has seen the release of numerous open-access LLMs, though
they generally do not perform as well as their closed-access counterparts.
“Open-access LLM” refers to models with publicly available weights, but
transparency varies regarding training data and filtering techniques.
EleutherAI released GPT-NeoX-20B[15] and GPT-J-6B [16], as well as the dataset
these models were trained on Pile[17]. Google released UL2-20B[18], an
encoder-decoder model trained on the publicly available C4[19]. Tsinghua
University released the weights of GLM-130B[20], a Chinese-English LLM, and
CodeGeeX-13B[21], a LLM for coding applications, without releasing the
training sets. Salesforce released CodeGen-Mono-16B[22] without disclosing a
proprietary Python dataset. Meta released the OPT[23], LLaMA[24] under a non-
commercial license and only provided high-level details about the data
collection and filtering process.

Many attempts to fine-tune on specialized domain-specific datasets only use
Supervised Fine-Tuning (SFT) or Keep Pretraining (KP), such as Lawyer
LLaMA[25] or Huatuo[26]. However, after an in-depth analysis of the dataset
and experiments using tokenizers of multiple pretrained models, we found that
using pre-trained tokenizers on domain-specific datasets is very unwise.
Tokenizers in general fields will produce redundancy several times or even
dozens of times in the token space, resulting in higher training resource
consumption and extended model convergence time. By utilizing the IGOT
approach to customize the tokenizer, our experiments demonstrate that the
model achieves several-fold improvements in information representation.
Additionally, this method conserves training resources and reduces
fluctuations in the loss function during training, thereby enabling the model
to converge to a more optimal point. To sum up, we make the following
contributions:

  * •

Our IGOT method enhances the model’s capability for information expression by
analyzing downstream task data to identify the most effective specialized
tokens for customizing the tokenizer. This process effectively transitions the
model’s tokenizer from a general domain to a vertical domain. Through
mathematical methods and experiments, we analyze the reasons why IGOT enables
the LLM to outperform traditional SFT or KP in multiple downstream tasks. This
investigation elucidates the superior performance of IGOT across various
applications.

  * •

We show that the IGOT method can be a more universal approach for customizing
models for downstream tasks. IGOT enhances model performance and conserves
training resources by shifting the tokenizer’s representational capabilities
to specialized domains. From a mathematical perspective, we analyze why a
customized tokenizer is essential for domain-specific LLMs, including
information entropy and data compression considerations.

##  II Motivation

Research by Chip-Chat[27] suggests that open-source LLMs fine-tuned on
hardware design data (e.g., Verilog) outperform generic models for these
tasks. This approach also avoids security risks associated with sending
proprietary data to external LLMs. CyberRio[28] also showed a way to use GPT-4
to harden a chip.

During our experiments with LLMs in specialized domains, we observed unusual
behavior in the vectorization of sentences containing domain knowledge by the
original LLaMA2 tokenizer. For instance, as shown in Figure 1(a), we continued
pretraining the LLaMA2 using documentation from OpenLane[29], which is an
automated RTL to GDSII flow based on components including OpenROAD, Yosys,
Magic, Netgen, CVC, SPEF-Extractor, KLayout, and a number of custom scripts
for design exploration and optimization. Subsequently, when we prompted the
model to introduce OpenLane, we noted that the LLaMA2 tokenizer divided
“Introduce OpenLane, an EDA tool” into 13 tokens, with “Introduce” being split
into “Int”, “rodu” and “ce” and “OpenLane” fragmented into “Open”, “L”, and
“ane”—splits that rendered the tokens meaningless.

As we mentioned in the introduction section, LLMs during their pretraining
phase not only learn about contextual text vectorization and probability
distributions but also acquire knowledge of grammar and syntax. The
nonsensical splitting observed with the original tokenizer clearly results in
a decreased ability of the subsequent Transformer architecture to capture
information effectively. Consequently, we proposed our IGOT method, as
illustrated in Figure 1(b). This approach allows the tokenizer to process user
inputs into more meaningful divisions during the tokenization stage, thereby
enhancing the informational content of the sentences. This method ensures that
the tokenizer adapts to domain-specific vocabularies and concepts, thus
preserving the semantic integrity and enhancing the model’s overall
performance in specialized tasks.

##  III Method

In this section, we will explain the implementation of our IGOT method, and
its heuristic optimization $IGOT_{\tau}$. By maximizing information gain, the
IGOT approach improves tokenization and enhances the model’s ability to
capture and utilize domain-specific knowledge more effectively.

###  III-A IGOT and $IGOT_{\tau}$

In the context of classification systems, the information gain brought by a
feature indicates its importance, with higher information gain suggesting
greater significance. A basic formula for conditional entropy is given by:

$H(Y|X)=-\sum_{x\in X}\sum_{y\in Y}P(x,y)\log P(y|x)$ (1)

In the application of LLMs, our input undergoes tokenization by the tokenizer,
resulting in a token list, which is then fed into the transformer architecture
to generate probability distributions for subsequent vocabulary. Assuming the
Transformer architecture can capture a context size of $\alpha$, with each
token represented by an 8-bit integer (capable of representing up to 65536
different tokens), if a meaningful word $\delta$ is split into $N$ tokens, and
a context consists of $\Delta$ words, we can define the unit of information
obtained by the Transformer’s window for this round as $\frac{N}{8\alpha}$. In
an ideal scenario, our IGOT would process the word $\delta$ as one token.
Thus, we can define the information gain $\theta$ of the idealized IGOT over
the entire context as:

$\theta=log(1+\sum_{i\in\Delta}N_{\delta_{i}})-log(1+\alpha)$ (2)

Under this premise, by setting a threshold $\epsilon$, we can simply construct
a meaningful set of words $\lambda$ from the downstream task:

$\lambda=\\{\theta_{n}>\epsilon,\ n\in\Delta\\}$ (3)

However, using $\lambda$ selected directly based on information gain might not
be the most efficient. Taking the OpenLane documentation as an example again,
the $\lambda$ obtained by running the IGOT method contains proprietary names
of chip IPs like “sky130A_sky130_fd_sc_hd_config” which are typically very
long and engineers often use a variety of custom symbols or abbreviations to
distinguish their meanings. This causes such words to be split into many small
tokens by the original tokenizer, resulting in very high information gain.
Meanwhile, effective words that we actually hope to capture, such as
“Tritonroute” may not achieve as high information gain, which is clearly not
desirable. To address this issue, we introduce a new heuristic method $\phi$,
which evaluates the information gain of words, their length, and the complete
downstream vocabulary to produce a score. Based on this score, we then obtain
a new subset of $\lambda$, denoted as $S_{\lambda}$:

$S_{\lambda}=\\{\phi(\theta_{n};\ n;\
\Delta)>\epsilon^{{}^{\prime}},n\in\Delta\\}$ (4)

We designate this heuristic-based method as $IGOT_{\tau}$.

###  III-B Heuristic $\phi$

To obtain the optimal $\phi$ function, we use a simple neural network trained
on a pre-annotated dataset D to develop an evaluator. This dataset D consists
of common words and a large number of words like
“sky130A_sky130_fd_sc_hd_config” Each word is scored from 1 to 5 to indicate
its desirability as a target word for our application. A simple idea is to use
a multi-class optimization objective with cross-entropy as the loss function
for optimization:

$\mathcal{L}(\Theta)=-\frac{1}{N}\sum_{i=1}^{N}\sum_{k=1}^{K}y_{ik}\log
p_{ik}+\lambda\sum_{j}\theta_{j}^{2}$ (5)

However, we prefer our evaluator $\phi$ to provide scores rather than
classifications. Therefore, we change our optimization objective to a
regression task, aiming to minimize the L2 norm between predicted values and
actual scores:

$\mathcal{L}(\Theta)=\frac{1}{N}\sum_{i=1}^{N}(y_{i}-f(x_{i};\Theta))^{2}+\lambda\sum_{j}\theta_{j}^{2}$
(6)

Training for multiple epochs with the Adam[30] optimizer yields the
$IGOT_{\tau}$ evaluator $\phi$.

###  III-C Domain Adaptive Pretraining

In our investigation, we focused on the LLaMA2 model within the framework of
the standard Causal Language Modeling (CLM) task. In this task, given an input
token sequence $x=(x_{0},x_{1},x_{2},\ldots)$, the model is trained to predict
the next token $x_{i}$ in an autoregressive manner. Mathematically, the
objective is to minimize the following negative log-likelihood:

$\mathcal{L}_{\textrm{CLM}}(\Theta)=\mathbb{E}_{x\sim\mathcal{D}_{\textrm{PT}}}\left[-\sum_{i}\log
p(x_{i}|x_{0},\ldots,x_{i-1};\Theta)\right]$ (7)

Here, $\Theta$ represents the model parameters, $\mathcal{D}_{\textrm{PT}}$ is
the pre-training dataset, $x_{i}$ is the token to be predicted, and
$x_{0},x_{1},\ldots,x_{i-1}$ constitute the context.

To our Domain Adaptive Pretraining (DAP) target, we will mask part of the full
task as the output, the loss is only calculated on the output part of the full
sequence and can be expressed as:

$\mathcal{L}_{\textrm{DAP}}(\Theta)=\mathbb{E}_{x\sim\mathcal{D}_{\textrm{DAP}}}\left[-\sum_{i\in\epsilon}\log
p(x_{i}|x_{0},\ldots,x_{i-1};\Theta)\right]$ (8)

Here, $\Theta$ represents the model parameters, $\mathcal{D}_{\textrm{DAP}}$
is the Domain Adaptation Pretraining dataset, $\epsilon$ is the output
sequence, $x=(x_{0},x_{1},\ldots)$ represents the tokenized input sequence.

##  IV Experiment

In Section 4.1, we begin by evaluating the impact of the IGOT method on the
convergence of training loss and discussing the reasons behind this
phenomenon. In Section 4.2, we generate heatmaps to analyze word distribution
with information gain that results in effective information gain for specific
datasets.

###  IV-A IGOT for DAP

(a) DAP with Original LLaMA2 Tokenizer

(b) DAP with IGOT

Figure 2: In Figure (a), we showed the performance of the original LLaMA2-7B
tokenizer during training. It is evident that even after the training reaches
the third epoch, the loss oscillation remains severe. However, Figure (b)
demonstrates that the IGOT method significantly reduces this oscillation,
facilitating better model convergence.

We conducted the Domain Adaptive Pretraining (DAP) on LLaMA2-7B using OpenLane
documentation and implemented the IGOT method. Figure 2 illustrates the impact
of the IGOT method on the training loss convergence. Without the utilization
of the IGOT method, the model’s convergence radius remains relatively large,
indicating substantial oscillation during training. This persistent
oscillation can hinder the model’s ability to converge to an optimal solution,
leading to suboptimal performance on downstream tasks.

Based on the results presented in Table I, we achieved significant savings in
computing resources, using IGOT with LLaMA2-7B on DAP leads to a 12.2% saving
in training time and 5% savings in GPU VRAM. On the T5[19] series model, we
were even able to achieve a 31.5% time-saving. On Meta’s OPT-350M model, there
are also varying degrees of time and GPU VRAM savings. The efficacy of the
IGOT method can be attributed to the customized tokenizer’s ability to capture
and incorporate domain-specific knowledge more effectively into the model’s
training process. By selecting optimal subsets of tokens based on their
information gain, the IGOT method ensures that the tokenizer adapts to the
nuances of the domain, resulting in more informative token representations. As
a result, the model can better leverage domain-specific information, leading
to improved performance on downstream tasks.

TABLE I: Computation savings brought by IGOT Approach | Added | Time Elapsed (s) | Max VRAM Usage
---|---|---|---
LLama2-7B | 0 | 4799.43 | 154GB
with $IGOT$ | 5232 | 4210.57 (-12.27%) | 146GB
with $IGOT_{\tau}$ | 640 | 4157.11 (-13.38%) | 147GB
OPT-350M | 0 | 86.34 | 22GB
with $IGOT$ | 3649 | 84.66 (-1.95%) | 21GB
with $IGOT_{\tau}$ | 619 | 80.09 (-7.23%) | 21GB
T5-Small | 0 | 292.88 | 14.5GB
with $IGOT$ | 4334 | 180.12 (-38.5%) | 13GB
with $IGOT_{\tau}$ | 615 | 241.45 (-17.56%) | 13GB
T5-Base | 0 | 301.16 | 28.7GB
with $IGOT$ | 4334 | 195.28 (-35.15%) | 27GB
with $IGOT_{\tau}$ | 615 | 251.02 (-16.64%) | 27.5GB

Furthermore, the use of a customized tokenizer also leads to more efficient
training, as evidenced by the observed reduction in training time and
resources. The more compact vectorization of sentences achieved by the
customized tokenizer allows for faster computation and more efficient
utilization of computational resources. This is particularly advantageous for
large-scale model training, where even minor improvements in training
efficiency can translate into significant time and cost savings.

###  IV-B Token Distribution

(a) $IGOT$

(b) $IGOT\tau$

Figure 3: Figure (a) depicts the distribution of information gain captured by
IGOT on the complete dataset. Although many tokens with information gain
greater than 4 are captured, the proportion is not high. Meanwhile, Figure (b)
shows that the supervised $IGOT_{\tau}$ exhibits a more extensive distribution
on high-gain tokens. This may be one of the reasons why the effectiveness of
$IGOT_{\tau}$ is somewhat better than IGOT.

The analysis shown in Figure 3 provides insight regarding the distribution of
tokens when using the IGOT method with OpenLane documentation. In Figure (a),
we can see the distribution of information gain captured by IGOT across the
entire dataset. While IGOT identifies many tokens with information gain higher
than 4, the proportion of such tokens is relatively low. This suggests that
although IGOT effectively identifies tokens with significant information gain,
there is still room for improvement in capturing a higher proportion of high-
gain tokens.

In contrast, Figure (b) illustrates the distribution of tokens when using
supervised $IGOT_{\tau}$. Here, we see a more extensive distribution of high-
gain tokens compared to the unsupervised IGOT method. This broader
distribution indicates that supervised $IGOT_{\tau}$ is more capable of
identifying and incorporating tokens with significant information gain into
the tokenizer. Consequently, the enhanced coverage of high-gain tokens in
supervised $IGOT_{\tau}$ may contribute to its superior effectiveness compared
to the unsupervised IGOT method.

Future research could focus on refining the criteria used for token selection
within the IGOT method. By developing more sophisticated algorithms or
heuristics for identifying tokens with significant information gain, we can
potentially enhance the effectiveness of the IGOT approach in capturing
domain-specific knowledge more comprehensively.

Another area for future exploration is the investigation of alternative
approaches to supervised token selection in the $IGOT_{\tau}$ variant. While
manual selection of token subsets based on information gain can yield
improvements, exploring automated or semi-automated methods for supervised
token selection could further streamline the process and potentially lead to
even greater gains in model performance.

Additionally, future work could involve the integration of domain-specific
constraints or knowledge into the tokenization process. By incorporating
domain-specific rules or guidelines into the tokenizer, we can ensure that the
tokenization process aligns more closely with the linguistic characteristics
and semantic nuances of the target domain, further enhancing the model’s
ability to capture and utilize domain-specific information effectively.

##  V Conclusion

This study proposed and validated the Information Gain Optimized Tokenizer
(IGOT) approach as a novel and effective strategy for adapting Large Language
Models (LLMs) to specialized domains. Our IGOT method significantly enhances
the domain-adaptation capabilities of LLMs by developing a customized
tokenizer that optimally utilizes domain-specific token sets determined
through heuristic function $\phi$ and their associated information gains.

Through our experimentation, including the use of LLaMA-7B for continued
pretraining on specialized datasets such as OpenLane documentation, IGOT
demonstrated a considerable improvement in model efficiency and effectiveness.
Specifically, our findings showed both more than 10% reduction in token usage
and training time, with a 5.8% reduction in maximum GPU VRAM usage. Such
enhancements not only streamline the training process but also optimize
computational resource utilization, making the adoption of general generative
AI technologies in specialized domains more feasible and efficient.

Our work with IGOT also highlighted the capacity of this method to retain
critical domain-specific information that traditional tokenization methods
often overlook or mishandle. By enabling more meaningful and semantically
coherent tokenization, IGOT not only enhances the quality of the model’s
output but also its ability to learn and apply new knowledge in domain-
specific contexts more effectively.

## References

  * [1] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al., “Language models are few-shot learners,” Advances in neural information processing systems, vol. 33, pp. 1877–1901, 2020.
  * [2] J. Vig, “A multiscale visualization of attention in the transformer model,” arXiv preprint arXiv:1906.05714, 2019.
  * [3] J. Hewitt and C. D. Manning, “A structural probe for finding syntax in word representations,” in Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4129–4138, 2019.
  * [4] C. Wang, X. Liu, and D. Song, “Language models are open knowledge graphs,” arXiv preprint arXiv:2010.11967, 2020.
  * [5] J. Davison, J. Feldman, and A. M. Rush, “Commonsense knowledge mining from pretrained models,” in Proceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference on natural language processing (EMNLP-IJCNLP), pp. 1173–1178, 2019.
  * [6] S. Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar, P. Lee, Y. T. Lee, Y. Li, S. Lundberg, et al., “Sparks of artificial general intelligence: Early experiments with gpt-4,” arXiv preprint arXiv:2303.12712, 2023.
  * [7] S. Wu, O. Irsoy, S. Lu, V. Dabravolski, M. Dredze, S. Gehrmann, P. Kambadur, D. Rosenberg, and G. Mann, “Bloomberggpt: A large language model for finance,” arXiv preprint arXiv:2303.17564, 2023.
  * [8] E. Bolton, A. Venigalla, M. Yasunaga, D. Hall, B. Xiong, T. Lee, R. Daneshjou, J. Frankle, P. Liang, M. Carbin, et al., “Biomedlm: A 2.7 b parameter language model trained on biomedical text,” arXiv preprint arXiv:2403.18421, 2024.
  * [9] K. Heafield, I. Pouzyrevsky, J. H. Clark, and P. Koehn, “Scalable modified kneser-ney language model estimation,” in Proceedings of the 51st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pp. 690–696, 2013.
  * [10] C. Buck, K. Heafield, and B. Van Ooyen, “N-gram counts and language models from the common crawl.,” in LREC, vol. 2, p. 4, 2014.
  * [11] Y. Bengio, R. Ducharme, and P. Vincent, “A neural probabilistic language model,” Advances in neural information processing systems, vol. 13, 2000.
  * [12] T. Mikolov, M. Karafiát, L. Burget, J. Cernockỳ, and S. Khudanpur, “Recurrent neural network based language model.,” in Interspeech, vol. 2, pp. 1045–1048, Makuhari, 2010.
  * [13] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” Advances in neural information processing systems, vol. 30, 2017.
  * [14] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al., “Language models are unsupervised multitask learners,” OpenAI blog, vol. 1, no. 8, p. 9, 2019.
  * [15] S. Black, S. Biderman, E. Hallahan, Q. Anthony, L. Gao, L. Golding, H. He, C. Leahy, K. McDonell, J. Phang, et al., “Gpt-neox-20b: An open-source autoregressive language model,” arXiv preprint arXiv:2204.06745, 2022.
  * [16] Y. Wang, W. Wang, S. Joty, and S. C. Hoi, “Codet5: Identifier-aware unified pre-trained encoder-decoder models for code understanding and generation,” arXiv preprint arXiv:2109.00859, 2021.
  * [17] L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, et al., “The pile: An 800gb dataset of diverse text for language modeling,” arXiv preprint arXiv:2101.00027, 2020.
  * [18] Y. Tay, M. Dehghani, V. Q. Tran, X. Garcia, J. Wei, X. Wang, H. W. Chung, S. Shakeri, D. Bahri, T. Schuster, et al., “Ul2: Unifying language learning paradigms,” arXiv preprint arXiv:2205.05131, 2022.
  * [19] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu, “Exploring the limits of transfer learning with a unified text-to-text transformer,” Journal of machine learning research, vol. 21, no. 140, pp. 1–67, 2020.
  * [20] A. Zeng, X. Liu, Z. Du, Z. Wang, H. Lai, M. Ding, Z. Yang, Y. Xu, W. Zheng, X. Xia, et al., “Glm-130b: An open bilingual pre-trained model,” arXiv preprint arXiv:2210.02414, 2022.
  * [21] Q. Zheng, X. Xia, X. Zou, Y. Dong, S. Wang, Y. Xue, Z. Wang, L. Shen, A. Wang, Y. Li, et al., “Codegeex: A pre-trained model for code generation with multilingual evaluations on humaneval-x,” arXiv preprint arXiv:2303.17568, 2023.
  * [22] E. Nijkamp, B. Pang, H. Hayashi, L. Tu, H. Wang, Y. Zhou, S. Savarese, and C. Xiong, “Codegen: An open large language model for code with multi-turn program synthesis,” arXiv preprint arXiv:2203.13474, 2022.
  * [23] S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, et al., “Opt: Open pre-trained transformer language models,” arXiv preprint arXiv:2205.01068, 2022.
  * [24] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al., “Llama: Open and efficient foundation language models,” arXiv preprint arXiv:2302.13971, 2023.
  * [25] Q. Huang, M. Tao, Z. An, C. Zhang, C. Jiang, Z. Chen, Z. Wu, and Y. Feng, “Lawyer llama technical report,” arXiv preprint arXiv:2305.15062, 2023.
  * [26] H. Wang, C. Liu, N. Xi, Z. Qiang, S. Zhao, B. Qin, and T. Liu, “Huatuo: Tuning llama model with chinese medical knowledge,” arXiv preprint arXiv:2304.06975, 2023.
  * [27] J. Blocklove, S. Garg, R. Karri, and H. Pearce, “Chip-chat: Challenges and opportunities in conversational hardware design,” in 2023 ACM/IEEE 5th Workshop on Machine Learning for CAD (MLCAD), pp. 1–6, IEEE, 2023.
  * [28] G. Y. Xinze Wang and Y. Zhu, “Cyberrio: A next generation risc-v agile hardware design using large language models.” https://github.com/hello-eternity/Cyberrio, 2024.
  * [29] M. Shalan and T. Edwards, “Building openlane: A 130nm openroad-based tapeout- proven flow : Invited paper,” in 2020 IEEE/ACM International Conference On Computer Aided Design (ICCAD), pp. 1–6, 2020.
  * [30] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” arXiv preprint arXiv:1412.6980, 2014.