Paper Title: Insurance Portfolio Pursuit with Reinforcement Learning.



The Inevitable Quark Three-Body Force and its Implications for Exotic States
    Su Houng Lee
    August 5, 2024
============================================================================



§ ABSTRACT 
	When faced with a new customer, many factors contribute to an insurance firm's decision of what offer to make to that customer. In addition to the expected cost of providing the insurance, the firm must consider the other offers likely to be made to the customer, and how sensitive the customer is to differences in price. Moreover, firms often target a specific portfolio of customers that could depend on, age, location, and occupation. Given such a target portfolio, firms may choose to modulate an individual customer's offer based on whether the firm desires the customer within their portfolio. Given a target portfolio, we term the problem of modulating offers to achieve this target portfolio the portfolio pursuit problem. We give a formulation of portfolio pursuit as a sequential decision making problem, and devise a novel reinforcement learning algorithm for its solution. We test our method on a complex synthetic market environment, and demonstrate that it outperforms a baseline method which mimics current industry approaches to portfolio pursuit. 




§ INTRODUCTION




 §.§ Background


Competition among the commercial insurance market has been intensified by the emergence of price comparison websites (PCW) as a core means by which customers and insurance firms interact. When a customer's details are entered into the website, many firms simultaneously decide what offers to make to the customer on the basis of those details. Since the customer can directly compare competing offers, firms must carefully calibrate their offers to maximise a trade-off between potential profit and the probability of acceptance. While both of these aims are local to each individual customer, firms may also wish to modulate the offers they make in order to pursue longer-term goals, realised over many customers. For example, the firm may wish to maintain a certain frequency with which new offers are accepted, to ensure they have the reserves necessary to pay out claims and avoid liquidity issues <cit.>. Alternatively, a firm may wish to break into a new segment of the market, and be content with making lower-than-optimal offers in order to attract a particular customer base. Goals such as these define a higher-level strategy that is crucial to the long-term success of the firm. Decisions regarding content of the long-term goals - the What? of the strategy - are typically made at a managerial level, and involve complex domain-expert knowledge and understanding of current market conditions. At present, the implementation of the pursuit of the goals - the How? of the strategy - also makes use of domain-expect knowledge. In this paper, we present a data-driven machine learning methodology for implementing higher-level strategies in the context of consumer insurance markets. 

The problem of what offer to make to a customer on a PCW can be divided into two sub-problems: the cost-estimation problem and the bidding problem. The cost-estimation problem can be summarised as the problem of estimating the total cost of insuring a given customer from the features available through the PCW. This is a well-studied supervised-learning problem <cit.>. We will not provide any further comment on the cost-estimation problem in the remainder of this paper; we take the solution to this problem as given. Instead, we focus on the bidding problem - what offer ought we to make to a customer, given an estimate of the cost of insuring them? The optimal price to offer a customer depends not only on the estimated cost of insuring them, but additionally on, (1) the offers that other insurers on the PCW will make to the customer; (2) the customer's sensitivity to price differentials; and, (3) whether the customer is the type of client the insurance firm wishes to insure. Our work focuses on this third aspect the bidding problem. 

Throughout our paper, we use the term portfolio to refer to the set of customers who accept an insurance firm's offers over some specified period of time (for example a day, week, or month). The problem of portfolio optimisation is that of deciding what the ideal portfolio would look like; how many of each different type of customer the insurance firm wishes to have concurrently insured. This may depend on aspects such as the kind of brand and reputation the firm wishes to cultivate, the desire to maintain a diverse customer base so as to protect against correlated risk or systematic model inaccuracies, or the necessity to avoid insuring too many customers and being unable to settle all the incoming claims. We take the firm's solution to the problem of portfolio optimisation also as given; that is, we suppose that the insurance firm already has a specification for the distribution of customers it would like to insure over some relevant future time-period. However, unlike in the stock market, where an investor can easily achieve their desired portfolio of stocks by simply investing and divesting, an insurance firm cannot arbitrarily add and remove customers from its portfolio. Rather, the firm must wait until a desired customer is asking for a quote, and then modulate the offer made to the customer in order to increase the likelihood that the offer is accepted. In deciding how to modulate the offer, the firm must carefully balance the value of having the customer in the portfolio against the profit they expect to make from insuring the customer. It would not be effective to, for example, make desired customers offers at which the firm is, on average, guaranteed a substantial loss. We term the problem of how to modulate offers made to customers in order to achieve a desired portfolio the problem of portfolio pursuit. 

It is to the problem of portfolio pursuit to which we make our contribution. Note that – unlike the cost-estimation, bidding, and portfolio optimisation problems – portfolio pursuit is intrinsically sequential in nature. If we wish to insure no more than 500 customers over the next week, the offers we ought to make to customers on Sunday depend heavily on our success rate for Monday through Saturday. As such, we cast the portfolio pursuit problem as a Markov Decision Process (MDP), and derive a novel Reinforcement Learning (RL) algorithm for its solution. Our decomposition into the four sub-problems given above - cost-estimation, pricing, portfolio optimisation, and portfolio pursuit - additionally enables us to develop a methodology which is minimally disruptive, in the sense that it can be integrated naturally with arbitrary solutions to the other three sub-problems. Indeed, our algorithm requires only a slight adjustment to existing methodologies for the bidding problem. 



 §.§ Previous work


To our knowledge, our paper is the first work to articulate a general mathematical formalism of the portfolio pursuit problem in insurance as a sequential decision making problem, and attempt to address it using reinforcement learning. <cit.> and <cit.> both apply reinforcement learning to other aspects of bid optimisation in the context of insurance. Like us, <cit.> considers the new business market for insurance. They have no notion of a desired portfolio, thus making each decision purely local. As such, they treat the problem as a contextual bandit, with no sequential aspect (beyond learning). In contrast, <cit.> examine the problem of pricing at renewal. While they consider the customer portfolio (thereby making the problem sequential), they restrict themselves only to constraints on the size of the portfolio, formulating it as a constrained Markov Decision Process. By additionally factoring in the internal structure of customer portfolios, our framework presents a more general and flexible approach to portfolio pursuit. 

Many previous works have attempted to apply reinforcement learning to the portfolio optimisation problem in the context of other markets, such as the stock market <cit.> and the cryptocurrency market <cit.>.  However, as noted above, this situation is fundamentally different from the case considered here. In these other markets, the relevant financial products can typically be purchased very quickly, allowing one to shift funds between them more or less freely (subject to transaction costs). This is in contrast to customers, who arrive sequentially and over a prolonged period. Furthermore, while purchasing cryptocurrency or a stock is a fairly simple transaction, bidding for a customer is not, since it involves attempting to out-compete alternative offers while accounting for stochasticity in a customer's decision. As such, this previous literature is not directly applicable to the problem under consider here. 



 §.§ Our contribution


We propose a novel reinforcement learning algorithm for the portfolio pursuit problem over long time horizons. We cast this as an MDP, and give a novel algorithm for its solution. We take a model-based approach, allowing our algorithm to be trained entirely offline, thereby protecting against risk in deployment. The methodology we use can be integrated naturally into existing machine learning pipelines for insurance pricing. Unlike previous research addressing using RL for the purposes of insurance pricing <cit.>, our insurance agent acts in a realistic agent-based environment <cit.>. This environment includes randomly generated customers along with a collection of other insurance agents that iteratively change their pricing strategies over time, creating a dynamic, non-stationary environment. The full code implementing both the market environment and our reinforcement learning solution is available in the https://github.com/EdwardJamesYoung/RL-portfolio-pursuitcode repository accompanying this paper. 

The rest of this paper is organised as follows. In sec:problem formulation we give a mathematical formulation of the portfolio pursuit problem. In sec:methods we describe our solution to the problem. We begin by describing a standard pipeline for solving the bidding problem, and the models used within that pipeline. We then outline a baseline methodology intended to mimic a naive approach to portfolio pursuit that may currently be used within industry. Lastly, we specify our approach, giving the full algorithm used for training, and compare the method to existing RL algorithms. In sec:results we show the performance of our method against the baseline, demonstrating that our method is able to generate more profit than the baseline methodology while not compromising the quality of the final portfolio. Finally, in sec:discussion we outline some limitations of our method and give opportunities for future research. 



§ PROBLEM FORMULATION




 §.§ Portfolio pursuit as a Markov Decision Process


In our market environment, an epoch consists of T interactions between a customer and the insurance agent. The customers s_1,…,s_T themselves are drawn independently from the same underlying customer distribution. At each step, in response to the current customer features, the insurer takes an action a_t ∈ [1,2]. The action a corresponds to an offer of C(s) × a, where C(s) is the expected cost of serving the customer[As noted in the introduction, we assume this as given. In practice, the expected cost of serving a customer is not known in advance, but must instead be estimated using historic claims data. However, this estimation problem is distinct from the problem of how to price insurance, given this estimate, which is what this paper focuses on.]. In response to the offer made by the insurer, and the offers made by other insurers within the market, the customer either accepts the offer or rejects it. If the offer is rejected (y_t = 0) then no profit is obtained by the insurer. If the offer is accepted (y_t = 1) then an (expected) profit of (s_t,a_t) is obtained. The profit is simply (s,a) = C(s)a - C(s) = C(s)(a - 1). When the offer is accepted (y_t = 1) then the customer is added to the insurer's portfolio[In reality, there may be a delay between an offer being made and accepted - indeed, customers often have up to 14 days to decide whether to accept an offer. In this case, a proxy measure, such as click-through from the PCW to the insurer's website, may be used in place.]. The portfolio at the t-th time-step is ρ_t, and the insurer's portfolio before the first customer ρ_1 is initialised as empty[The choice to initialise the portfolio as empty corresponds to optimising the portfolio of new customers obtained over some time horizon. If instead the firm is interested in its total portfolio - including all existing customers - the portfolio can be initialised to that instead, with only minor adjustments required elsewhere in our method.], ρ_1 = {}.

The performance of the insurer at the final time step is evaluated according to the net profit achieved minus a loss which measures the discrepancy between the final and desired portfolio[Note that this encompasses the case where the loss is a function of only the final portfolio.], 
    ∑_t=1^T y_t (s_t,a_t) - λŁ(ρ_T+1, ρ^*).

Here, λ is a coefficient which controls the relative importance of profit vs. pursuit of a particular portfolio. 

In the MDP formalism, the states consist of customer, portfolio, and time-step triples, (s_t, ρ_t, t). The reward signal following action a_t is y_t (s_t,a_t), with an additional reward -λŁ(ρ_T+1, ρ^*) administered after the final customer interaction.



 §.§ Specification of portfolio loss function


We stress that the algorithm we present in sec:RL for PP makes no assumptions whatsoever about the nature of the loss function Ł. However, for concreteness in our numerical experiments, we adopt one particular formulation of the loss. 

The loss function used in our numerical experiment is defined via a collection of I indicators functions _A_i, i = 1,…,I over the customer features. Each indicator is specified by set A_i, i = 1,…,I in the space of customer features. These indicators allow us to represent portfolios as vectors in ^I, with the i-th component given by f_i = ∑_s ∈ρ_A_i(s), where _A_i(s) is equal to 1 if s ∈ A_i, and 0 otherwise. Given these vector representations of portfolios, we can implement the loss Ł via any distance measure on ^I. Note that in practice we specify the target portfolio ρ^* only through its vector representation f^*_i ∑_s ∈ρ^*_A_i(s); similarly, we will track the insurer's portfolio only via its frequency representation. For our numerical experiments we adopt the loss function:

    ℒ(f, f^*) 1/I∑_i=1^I   |f_i - f^*_i| /max(f_i, f_i^*) .

We choose this loss for the following properties: it is zero precisely when f_i = f^*_i; it is bounded between [0,1]; and it is symmetric in f and f^*. Note that because the loss is bounded between 0 and 1, the coefficient λ can be interpreted as answering the question: if the firm had the worst possible portfolio, how much short-term profit would they be willing to give up in order to trade it for the best possible portfolio? 

Our reinforcement learning algorithm (detailed in sec:RL for PP) trivially extends to alternative loss functions which provide other ways of quantifying the distance between the obtained and target portfolio. This includes, for example, adding weights to each indicator to penalise discrepancies along different feature sets differently, or using an optimal transport distance to provide a more detailed measure of differences between sets. We leave exploration of alternative loss formulations to future work. 



§ METHODS




 §.§ Standard industry methodology


We compare our method to a pipeline which mimics current industry approaches to insurance pricing and price modulation. In this section we describe a methodology which does not make reference to a target portfolio. This pipeline mimics standard industry pricing methods. We will then modify this method in two ways - firstly, to mimic standard industry methodology for modulating prices to attain a target portfolio, and secondly according to our new proposed algorithm.  

The insurance firm is assumed to have a dataset of interactions from previous epochs. Each entry in this interaction dataset includes:

	
  * The customer features, s_t
	
  * The expected cost of serving the customer, C(s_t)
	
  * The action taken by the insurance firm, a_t
	
  * The offers made by other insurance firms, o_t
	
  * Whether the customer accepted the offer of the firm, y_t

From the offers made by other firms, o_t, and the expected cost, C(s_t), the insurer calculates market variables, m_t. The market variables for the insurer are the average of the top-1, top-3, and top-5 offers made to the customer by other firms within the market, normalised by dividing by the expected cost C(s_t) of serving the customer. Note that the market variables are not available to the insurer at inference time, but become available after the fact through the purchasing of historic data from PCWs. 

From this dataset, the insurer trains the following models:

	
  * The market model M, which maps from customer features, s_t, to market variables m_t, s_t ↦ M(s_t) ≈ m_t.
	
  * The conversion model p, which maps from the market variables, m_t, and action, a_t, to the probability the customer accepts the offer, (m_t,a_t) ↦ p(m_t,a_t) ≈(y_t = 1|m_t, a_t).
	
  * The action model, or policy π, which maps from market variables, m_t, to optimal actions a^*_t, m_t ↦π(m_t) ≈ a^*_t.

For our numerical simulations, we used a random forest for the market model, a Multi-Layer Perceptron (MLP) for the conversion model, and a Gaussian Process with Matern kernel for the action models <cit.>. Our overall methodology places no restrictions on any of the models, although in practice we found it better to use continuous conversion and action models. If it is necessary to ensure these models are interpretable, simpler models - such as Generalised Linear Models (GLMs) - may be used instead of MLPs <cit.>. More expressive models may work better for larger scale simulations, with correspondingly larger datasets, than those explored here. 

The market model M and conversion model p are trained directly on the historic dataset available to the insurer. To train the action model, we use regression-based amortised optimisation <cit.>. Firstly, note that the expected profit, given that the customer accepts the insurer's offer, is C(s_t)(a_t - 1). Accordingly, the insurer's expected profit, for action a_t, given market variables m_t, is C(s_t)(a_t - 1)p(m_t, a_t). Therefore, for a random sample of market variables m_t from the dataset, we use a one-dimensional optimiser to find a^*_t ∈max_a p(m_t, a)(a - 1) within the range a_t ∈ [1, 2]. We then fit the action model π to 500 (m_t, a^*_t) pairs. 

At inference time, the features for a customer s_t are fed into a market model to obtain estimates for the market variables[As the market variables are predictive of customer behaviour, and available historically but not at inference time, the use of a market model can be seen as a form of data imputation.], m_t = M(s_t). These are then fed into the action model to obtain the action a_t = π(m_t). An offer of C(s_t)a_t is then made to the customer. On 12% of customers we explore by multiplying the offer by a randomly chosen modulation factor sampled from [0.93, 0.95, 0.97, 1.03, 1.05, 1.07]. 

Before continuing, we briefly remark that in practice the training and use of auxiliary models may be more complicated than presented here. To begin with, information about whether a customer accepts an offer is not available instantaneously, but only after a “cooling-off” period, typically 14 days. Likewise, there may be a delay in obtaining information about market variables for a given customer. Finally, insurers may choose to forgo the use of an action model all-together, instead solving the one-dimensional optimisation problem directly at inference time (rather than amortising as we do here). None of these observations present substantial difficulty to the RL algorithm we describe below. Indeed, our approach can be seen as a methodology which accepts as input auxiliary models and adjusts the prices returned by those models to solve the portfolio pursuit problem. It is therefore agnostic as to both the types of models used, and the data which those models are trained on. 



 §.§ Baseline methods for portfolio pursuit


It is standard practice to modulate the output prices offered to customers from the bidding model according to other features of interest for the customer, corresponding to the customer's desirability within the portfolio. This modulation is typically performed using domain experts' knowledge, rather than according to a data-driven or formal scheme. We attempt here to give a method which mimics this type of modulation as a baseline against which to compare our method, fully recognising that it does not accurately capture the complexity of standard industry practices. We encourage future work to implement more realistic and nuanced methodologies, and have facilitated this by making our codebase implementing the market environment publicly available. 

The baseline method makes use of historical assessments of the frequency with which the insurer encountered customers of the type relevant to their portfolio. In particular, for each indicator set A_i, compute the frequency representation f_i = ∑_s ∈ρ_T+1_A_i(s) for terminal portfolios ρ_T+1 for epochs during which there was no portfolio-based action modulation, where actions were chosen via the method described in the previous section, sec:standard industry method. We average this value over such historical epochs, giving us a frequency f̅_i. f̅_i therefore acts as an estimate of how many customers within set A_i we are likely to get without action modulation. We compare this to the desired frequency, f^*_i = ∑_s ∈ρ^*_A_i(s). For each incoming customer, we look at each set A_i that they call into, and compute the corresponding ratio of frequencies, f̅_i/f^*_i. If this ratio is greater than 1 it indicates that, in order to achieve the desired portfolio, we should increase the offer price made to that customer. If the ratio is less than 1, we should decrease the offer price. At inference time, we therefore modulate the offer price made to a customer s by multiplying the action by

    ∏_i: s ∈ A_i g( f̅_i/f^*_i).

The map g is given by 

    g(z) = 1 + βz^n - 1/z^n + 1,
 
where n, β≥ 0. The functional form of g was selected because it is monotonically increasing and bounded, with g(z) ∈ [1 - β, 1 + β]. The function parameters n ∈{0.5, 1, 2} and β∈{ 0.02, 0.05, 0.1, 0.25 } are found by a grid-search; in our experiments we used values of n = 1, β = 0.02 for a loss coefficient of λ = 2000. Note that for a baseline method, β, n must be chosen differently for each value of the loss coefficient λ. This is in contrast to our method, which applies equally well for any value of λ. 



 §.§ A Reinforcement Learning approach to Portfolio Pursuit




  §.§.§ Mathematical Preliminaries


Our reinforcement learning algorithm uses a dynamic programming approach with function approximation to provide a policy which modulates customer prices according to the current portfolio possessed by the insurer. We define the portfolio value function, V_π(ρ, t), which gives the value of having portfolio ρ at time t, given that future actions are selected by a (time-step dependent, deterministic) policy π = π(s, ρ, t):

    V_π(ρ, t) _π[ ∑_k=t^T y_k (s_k, a_k) - λŁ(ρ_T+1, ρ^*) | ρ_t = ρ].

The corresponding action value function Q_π(s,a,ρ,t) is defined by:

    Q_π(s, a, ρ, t) _π[ ∑_k=t^T y_k (s_k, a_k) - λŁ(ρ_T+1, ρ^*) | ρ_t = ρ, s_t = s, a_t = a ].

The policy is optimal for all time steps if it satisfies the appropriate Bellman optimality equation, 
    π(s, ρ, t) ∈max_a Q_π(s, a, ρ, t).

In app:bellman recursion, we show by expanding the action value function that (<ref>) is equivalent to:

    π(s, ρ, t) ∈max_a (y_t = 1 |s_t = s, a_t = a) ( a - k_π(s,ρ, t+1) ),
 
where k_π(s, ρ, t+1) is defined by:

    k_π(s, ρ, t)  1 - V_π(ρ∪{s}, t) - V_π(ρ, t)/C(s).

An important consequence of the above result is that the optimal policy is a function only of the customer features s_t and the value k_π, and not of the portfolio or time-step directly except via k_π.

We also obtain (see app:bellman recursion) the following recursion relationship for the portfolio value function:

    V_π(ρ, t) = V_π(ρ, t+1) + [  C(s) (y_t = 1 |s_t = s, a_t = π(s,ρ, t)) ( π(s,ρ, t) - k_π(s, ρ, t+1) )  ],

with boundary condition V_π(ρ, T+1) = - λŁ(ρ, ρ^*). Here, the expectation is taken over the customer distribution. 



  §.§.§ Outline of our method


As with the industry baseline method, our method will make use of a market, conversion, and action model. We will use the conversion model p(m, a) ≈(y_t = 1 | m_t = m, a_t = a) in place of the (y_t = 1|s_t = s, a_t = a) in both the value recursion relationship (<ref>) and the optimal action relationship (<ref>)[We stress that this is only an approximation. In reality – owing to how we constructed the market environment – the probability of acceptance also depends on various customer features.]. Additionally, our action model (or policy) π will now be a function of both the market variables m and the k-values K, π(m,k). 

Our method begins by training these models in much the same manner as the pipeline described in sec:standard industry method, namely by fitting models to a previously obtained dataset. However, our method for generating the dataset for fitting the action model π is slightly different. We first sample market variables m_t from our dataset. For each market variable sample, we independently sample a k-value, k_t, from a Laplace distribution with location 1 and scale 0.1. We then compute 

    a^*_t ∈max_a p(m_t, a)(a - k_t),
 where p is our conversion model. We then fit the action model to approximate π(m_t, k_t) ≈ a^*_t. We used a sample of 500 (m, k, a^*)-triples to fit our conversion model. We additionally train a value function V_ϕ(ρ, t), with parameters ϕ, as outlined in the next section. 

At inference time, for a customer s_t with expected cost C(s_t) and portfolio ρ_t at time t, we use (<ref>) to compute the k-value k_t via:

    k_t = 1 - V_ϕ(ρ_t∪{s_t},t+1) - V_ϕ(ρ_t,t+1)/C(s_t).

We then feed the customer features s_t into the market model to obtain estimates for the market variable m_t = M(s_t). We then use the bidding model to obtain the action a_t = π(m_t, k_t). Finally, we make an offer of C(s_t)a_t to the customer. 



  §.§.§ Value function training


In order to train our inference-time value function V_ϕ(ρ,t), we make use of an next-step value function, U(ρ). This function is used during training, but not at inference time. Throughout training, it approximates the values of portfolios at the next time step. U is used to create value estimates, which are then stored in a dataset. Our value function V_ϕ (which will be used at inference time) is then fit to this dataset. As discussed in sec:portfolio loss function, portfolios are specified only through their frequency representation, ρ↦ (f_i)_i=1^I, f_i = ∑_s ∈ρ_A_i(s), since this preserves all information relevant for computing the loss. Accordingly, both U and V_ϕ act on frequency representations, vectors in ^I. In the remainder of this section we do not distinguish between a portfolio and its frequency representation. 

We iterate backwards through time, setting t = T,…,1. At each time step, we use U(ρ) to approximate V_π(ρ, t+1). Initially, we set U(ρ) = -λℒ(ρ, ρ^*) = V_π(ρ, T+1). At each time step, we use the next-step value function U to form value estimates via (<ref>) and (<ref>). To approximate the expectation over the customer distribution in (<ref>), we compute an empirical average over samples from a customer reply buffer, which stores the features, market variables, and expected cost of previous customers. We then fit U to these value estimates thus obtained, such that U now approximates V_π(ρ, t). Having fit U, we sample an augmenting set of portfolios, and use U to estimate each of their values. We finally store the value estimates in a dataset, 𝒟, and then iterate backwards, t  t - 1. Once we have completed the backwards iteration through time, we fit our value function V_ϕ to the value estimates found in 𝒟.

Before fitting, we apply a trick to make our dataset easier to learn. Notice that at inference time, V_ϕ is used only to compute k values via (<ref>). However, this equation is invariant to shifting all values at same time step t by the same amount. Therefore, before fitting V_ϕ, we recentre all the value estimates V at each time step by subtracting the mean value estimate, V̅.  We then fit our value function V_ϕ to the residuals, V̂ = V - V̅. We found that this step is critical to effective learning, as it requires the value function V_ϕ to learn the differences in value between portfolios, rather than overall magnitude. These differences are what drives decision making - indeed, the centring of our value estimates makes V_ϕ analogous to an advantage function used in other RL methods. 

Our approach is summarised in alg:training value function. In our experiments, we set U to be linear model, and V_ϕ to be an MLP. As with the conversion, market, and action models, nothing in our algorithm necessitates the use of these particular models for U and V_ϕ. We use N = 500 samples from the customer reply buffer, and J = 24 portfolio samples for fitting the next-step model, and an additional J^+ = 120 augmenting portfolios. 



In our algorithm, there are two points at which we sample portfolios. We have not yet specified how this sampling occurs. Ideally, these portfolios would be on-policy, the portfolios sampled for time step t would be taken from the distribution of portfolios our policy encounters at time t. However, since we do not have access to any such data, our method is forced to be off-policy. The sampling of portfolios comes from three sources: previously on-policy portfolios, target on-policy portfolios, and high-coverage portfolios. The sampling is done in a ratio of 1:1:2. We explain each of these methods in turn. 

Much like the baseline method discussed in sec:baseline method, the previously on-policy portfolio sampling makes use of historic data about the frequency with which customers in each category appear naturally without action modulation. As before, we denote by f̅_i the average of the frequency representations ∑_s ∈ρ_t+1_A_i(s) over historic epochs in which no portfolio was being pursued. We can then define a historic rate, given by p̅_i = f̅_i/T. To sample previously on-policy portfolios at time t, we sample each component of the portfolio's frequency representation independently from the Binomial distribution with parameters p_i + 1/T, t. The addition of 1/T to the rate was made to ensure that, for each category, we have a strictly positive probability of sampling portfolios with customers in that category. Sampling these portfolios ensures we have a good representation of value for portfolios that we would encounter if we did not adjust our pricing methodology. The target on-policy portfolio sampling method works in the same manner as the previously on-policy portfolio sampling, except we use p^*_i = f^*_i/T where f^*_i = ∑_s ∈ρ^*_A_i(s) is the target frequency. 

The last method we use is high-coverage portfolios. This method is parametrised by a slope parameter, σ∈ [0,1]. We first define p_max = (1 + σ)max( p^*, p̅ ), p_min = (1 - σ)min( p^*, p̅ ), where the maximum and minimum are taken element-wise. At time t, we consider evenly spaced points ω∈ [0,1]. For each ω point, we construct the frequency representation f_t(ω) = t( ωp_max + (1 - ω)p_min ), which interpolates between the frequency representation if customers were acquired at either the maximum rate or minimum rate. This will in general not be integer valued, so we therefore take the integer part of each element, and then add on a Bernoulli sample which is one with probability equal to the remainder term. The high-coverage portfolios thus sampled ensures a wide sampling outside the on-policy distributions. The contribution of these portfolios is crucial, as it ensures our method still has accurate assessments of portfolio values outside the on-policy distribution, providing stability at inference time. We used a value of σ = 0.9.



 §.§ Relationship to existing RL methods


Here we briefly discuss the relationship of our proposed RL algorithm to existing algorithms within the literature. Our method can be viewed as an actor-critic algorithm, in the sense that it learns both a policy π(m,k) and a value function V_ϕ(ρ,t). However, our method displays substantial differences to more conventional actor-critic methods, such as A3C <cit.>, PPO <cit.>, SAC <cit.>, and DDPG <cit.>. 

Firstly, our method is not a policy-gradient algorithm. Unlike the methods listed above, our policy is not learned by finding the gradient of a value measure (such as total return <cit.>, or a Q-function <cit.>) with respect to the policy parameters. Instead, we exploit the fact that - because the action space is one-dimensional - the relevant value measure, given in (<ref>), can be optimised directly with a non-gradient-based optimiser. We then fit our policy directly to the optimised values - a form of regression-based amortised optimisation <cit.>. This allows us to avoid problems relating to local maxima which arise in gradient-based methods. It also frees us from the necessity of having a policy which is differentiable in the parameters; indeed, in our experiments we use a Gaussian process as our policy <cit.>. In fitting the policy to a dataset of actions which maximise a value measure, our method is akin to the cross-entropy method <cit.>. 

Secondly, our value function is not a function of the full MDP state. The full MDP state is the (customer, portfolio, time-step) triple, (s,ρ,t). Instead, our value function only depends on the portfolio and the time-step, (ρ,t). This substantially reduces the dimensionality of the input space to our value function, since the number of customer features is (typically) much larger than the number of portfolio categories. From (<ref>), we know that the optimal action depends only on the next portfolio, time-step pair. This means that, for the purposes of action-selection, it suffices to learn only these values, rather than the full MDP state values (or, indeed, the value of state-action pairs <cit.>). We also note that the dependence is only through the k-value. This enables us to learn a policy which takes k directly as input, rather than the full MDP state. Importantly, this policy can be learned fully before learning any value functions (as we do in our algorithm), and does not require re-training while the value function is learned, unlike a typical actor-critic method <cit.>. 

Lastly, our method utilises a memory-reply buffer to act as a model of the arrival of new customers. However, unlike many other methods which use a memory reply buffer for value-learning <cit.>, we do not sample states randomly. Rather, (portfolio, time-step) pairs, (ρ, t), are sampled systematically by iterating t backwards in time. In iterating backwards, our method is analogous to a dynamic programming algorithm, which starts at terminal states of known value (in our case, terminal portfolios) and applies backwards induction to find the value of all states recursively. Similar to the Fitted Q-Iteration algorithm <cit.>, our method avoids the curse of dimensionality associated with exact dynamic programming in MDPs by approximating the value function at each time-step by a linear model, U. Because we are applying an approximate dynamic programming method we need only sweep through the state-space once. This is in contrast to methods which apply a Bellman optimality equation but do not recurse backwards from terminal states <cit.>, where each state must be sampled many times in order to approximate the correct value. 





  §.§.§ Other applications of our algorithm


Although we have designed our algorithm specifically for the portfolio pursuit application pursued in this paper, certain elements readily generalise to other situations not considered here. We briefly enumerate potential other uses of elements of our algorithm. 

Firstly, our algorithm depends on the fact our state-space factorises into portfolio, time-step, customer feature triples (ρ, t, s). The fact that we have a fixed horizon T, and that the time-step t is a component of the state-space, allows us to compute values by iterating backwards in time. 

the use of a simple (in our case, linear) model 






 §.§ Additional simulation details


Our simulations use a realistic agent-based model of an insurance market, mediate through a PCW <cit.>. The insurance firm implementing each method – baseline (sec:baseline method) or RL (sec:RL for PP) – competes with 5 other insurance firms (none of which are pursuing a portfolio, as in sec:standard industry method). Each customer is randomly generated by sampling features from a complex generative distribution, with realistic interdependencies between features. For example, customer income depends on the location, age, and occupation of the customer. Only a subset of features are revealed to the insurers though the PCW interface. On the basis of these features, each insurer makes the customer an offer. The customer then either selects an offer or chooses to walk away from the market. Full details of the market environment can be found in the https://github.com/EdwardJamesYoung/RL-portfolio-pursuitcode repository that accompanies this paper. 

Target portfolios are generated at random by sampling I = 5 customer features, and setting the target quantity f_i^* for those features to be twice f̅_i if f̅_i ≤ 10, and otherwise half or twice of f̅_i at random (rounding up or down respectively to ensure integer values). We have 6 burn-in epochs, during which historic data is gathered for training models. Insurers use the past 4 epochs to train models, and all burn-in epochs for computing historical frequency representations. For the first burn-in epoch, before any models are trained, insurers take actions randomly and uniformly in [1,1.2]. Following the burn-in epochs, we then have 8 testing epochs. Each epoch consists of T = 1000 steps. No re-training of any models occurs following a testing epoch, meaning the testing epochs are independent trials of the same value function V_ϕ. In the results below, we aggregate statistics across testing epochs. We set the loss-coefficient λ = 2000. The results shown in sec:results are obtained by aggregating statistics across both the 8 testing epochs and 24 independent trials, with each trial having a different randomly generated target portfolios. To ensure a fair comparison between the RL and baseline methods, we use the same set of random seeds for both methods. This means that, up until the testing epochs, both methods interact with the market identically, ensuring the same historic data to train on and target portfolio to aim towards. Thus, any differences in performance are attributable only to differences over the testing epoch. 



§ RESULTS




fig:main-results shows the results of comparing our method to the baseline averaged over 24 trials. On each trial, both methods have the same interaction history, and thus have the same training data, up until the testing epochs. They also pursue identical target portfolios. Thus, differences in performance stem only from different behaviour on the test epochs. From panels (A-B) we see that, on average, our method generates more profit than the baseline (one-sided Mann-Whitney U test, p < 0.0001; CLES = 65%; Cohen's d = 0.51, moderate effect size). Over the course of 1000 customer interactions, on average our method generates an additional 472 (or a 7% increase) in profit. Panels (C-D) show the loss between the target portfolio and the obtained portfolio at each time step. The final loss distribution of both methods is very similar (two-sided Mann-Whitney U test, p = 0.16, not significant), showing that our method has not compromised portfolio quality in order to generate additional profit. Thus we can conclude that our method brings in additional revenue while getting comparably close to the desired portfolio as the baseline. Panels (E-F) show the total reward (profit minus loss) of both methods, which is the objective that both methods are optimising. We see that our method outperforms the baseline, achieving higher reward on average over the trials (one-sided Mann-Whitney U test, p < 0.0001; CLES = 63%; Cohen's d = 0.48, moderate effect size). In Table <ref> we tabulate the average final profit, loss, and reward of both methods.  





§ DISCUSSION




 §.§ Limitations and opportunities for future work


We believe that our work presents an important first step in the application of ideas from reinforcement learning to the portfolio pursuit problem. However, our work has several limitations, with corresponding opportunities for future research. 

In our investigation, we limited the time span over which insurers pursued a particular portfolio to only T = 1000 time-steps, with one customer arriving per time-step. This is a relatively short length of time for a large insurer, who may see orders of magnitude more customers per day. Our choice to limit the number of time-steps was motivated primarily by the prohibitive computational cost of running dozens of market simulations for larger numbers of time steps. Preliminary experiments provided evidence that our method will likely continue to work for larger numbers of time-steps. However, for much larger lengths of time, on the order of months, a different approach may well be needed. Accordingly, we believe that the framework and methods we have develop above could be integrated naturally with RL methods which attempt to perform temporal abstraction, goal-conditioned policies <cit.> or the options framework <cit.>. Indeed, longer time-scale portfolio pursuit can be accomplished by chaining together shorter time-scale portfolio pursuit with the appropriately chosen intermediate target portfolios. We leave the problem of selecting such intermediate target portfolios to future research. 

In our framework, we have also made the simplifying assumption that customers do not leave the portfolio once they have entered. This is justified when the number of time-steps is small, since few customers will both accept and then cancel a contract over the course of, one day. However, for longer time-scales, it may begin to become necessary to model the exit of customers from the portfolio. In app:customers leaving we expand our analytic derivation to consider this case, and provide some preliminary suggestions for how to integrate this into our algorithm. 

Lastly, our method is model-based, with optimal actions and values computed using a conversion model which estimates the probability of customers accepting any given offer. Consequently, if the conversion model is incorrect, our bidding strategy will accordingly be sub-optimal. For example, the bidding model may underestimate the probability that a customer accepts an offer at a given price, and make offers too low. Model uncertainty could be incorporated into the algorithm by explicitly penalising the value of actions about which the insurer is uncertain. A core step in making the algorithm more robust would be allowing for online updates of the models in light of new data coming in. This is an increasingly pressing problem within insurance pricing in particular, since such markets tend to be very non-stationary, shifting rapidly as individual insurers change their pricing strategy and inflation drives changes in prices. An effective pricing strategy will need to precisely trade-off between profit and the value of the information generated by making particular offers. Although we do not address it in this paper, we expect this to be a fruitful area for future research into the application of RL to insurance pricing. 



 §.§ Conclusion


Our paper has introduced a new formalism for a problem within the insurance industry - that of portfolio pursuit. Portfolio pursuit asks the question of how to modulate the offers made to individual consumers, in light of a desired target portfolio which describes the higher-level strategy of the insurance firm. We formulated portfolio pursuit it as an MDP in sec:problem formulation. Given its formal description, in sec:methods we discussed both industry standard methodologies for addressing this problem, and gave our own reinforcement learning solution. We then showed in sec:results that our method outperforms the industry baseline, generating more profit while getting equally close to the target portfolio. It is our hope that this paper will stimulate additional research into this area, laying the groundwork for better algorithms and methodologies for portfolio pursuit. 



 §.§ Acknowledgements


EJY is supported by the UKRI Engineering and Physical Sciences Research Council Doctoral Training Program grant EP/T517847/1. Work was completed during an internship at Accenture (UK) Limited, in connection with the Turing Institute. 



 §.§ Declaration of competing interest


The authors declare no competing interests.



 §.§ Codebase


The GitHub code repository associated with this paper can be found https://github.com/EdwardJamesYoung/RL-portfolio-pursuithere. The code is licensed under Creative Commons Attribution-NonCommercial 4.0 International Public License (CC BY-NC 4.0)



 §.§ Author contributions


Problem statement was formulated jointly by EJY, AR, and JJ. Market environment designed by EJY. Mathematical derivations, algorithm design, and numerical experiments were preformed by EJY. EJY drafted the manuscript, with additional comments and feedback provided by ET, AR, and JJ. 



apalike








§ BELLMAN RECURSIONS FOR PORTFOLIO VALUES


In this appendix we expand out the Bellman equations for the portfolio and action value functions, given by (<ref>) and (<ref>) respectively. We start with the action value function. We will expand this in terms of the portfolio value function. Firstly, we denote the returns by

    G_t = ∑_k=t^T y_k (s_k, a_k) - λŁ(ρ_T+1, ρ^*).

Then we have

    Q_π(s, a, ρ, t)    _π[ G_t | ρ_t = ρ, s_t = s, a_t = a ] 
       = (y_t = 1 |s_t = s, a_t = a)[ G_t | ρ_t = ρ, s_t = s, a_t = a, y_t = 1 ] 
       + (1 - (y_t = 1 |s_t = s, a_t = a))[ G_t | ρ_t = ρ, s_t+1 = s, a_t+1 = a, y_t = 0 ] 
       = (y_t = 1 |s_t = s, a_t = a)( (s,a) + [ G_t+1 | ρ_t+1 = ρ∪{ s }]  ) 
       + (1 - (y_t = 1 |s_t = s, a_t = a))[ G_t+1 | ρ_t+1 = ρ] 
       =(y_t = 1 |s_t = s, a_t = a)( (s,a) + V_π( ρ∪{ s }, t+1 ) ) 
       + (1 -(y_t = 1 |s_t = s, a_t = a)) V_π( ρ, t+1 ) 
       = V_π(ρ, t+1) + (y_t = 1 |s_t = s, a_t = a) ( (s,a) + V_π( ρ∪{ s }, t+1 ) -  V_π( ρ, t+1 ) ).

Note that, since the profit is the offer price C(s)a minus the expected cost C(s), the profit is thus given by: (s,a) = C(s)(a - 1). We can then say that:

    Q_π(s,a,ρ,t)    = V_π(ρ, t+1) + (y_t = 1 |s_t = s, a_t = a) ( C(s)(a - 1) + V_π( ρ∪{ s }, t+1 ) -  V_π( ρ, t+1 ) ) 
       = V_π(ρ, t+1) + C(s) (y_t = 1 |s_t = s, a_t = a) ( a - k_π(s,ρ, t+1) ),

where k_π(s, ρ, t+1) is defined in (<ref>). Since the first term is a constant in a, and C(s) > 0, this shows that maximisation of Q_π over a is equivalent to maximising:

    (y_t = 1 |s_t = s, a_t = a) ( a - k_π(s,ρ, t+1) ).
 
Furthermore, note that V_π(ρ, t) = [ Q_π(s,a,ρ, t) ]. Assuming a deterministic policy π = π(s,ρ, t), this gives us the following recursion:

    V_π(ρ, t) = V_π(ρ, t+1) + [  C(s) (y_t = 1 |s_t = s, a_t = π(s,ρ, t)) ( π(s,ρ, t) - k_π(s, ρ, t+1) )  ].




§ GENERALISATION TO CUSTOMERS LEAVING


In this section we generalise the analysis is app:bellman recursion to the case where customers can leave the portfolio. 

We define the action-value function as:

    Q_π(s, a, ρ, t) _π[ G_t | ρ_t = ρ, s_t = s, a_t = a ]

We will now compute the Bellman recursion for Q_π. To do so, we introduce a portfolio transition distribution, p( ρ' | ρ ). This gives the probability of having portfolio ρ' at time t+1, given that you ended time-step t with portfolio ρ. 

    Q_π(s, a, ρ, t)    = (y_t = 1 |s_t = s, a_t = a)[ G_t | ρ_t = ρ, s_t = s, a_t = a, y_t = 1 ] 
       + (1 - (y_t = 1 |s_t = s, a_t = a))[ G_t | ρ_t = ρ, s_t+1 = s, a_t+1 = a, y_t = 0 ] 
       = (y_t = 1 |s_t = s, a_t = a)( (s,a) + ∫[ G_t+1 | ρ_t+1 = ρ' ] p(ρ' | ρ∪{s} ) dρ' ) 
       + (1 - (y_t = 1 |s_t = s, a_t = a)) ∫[ G_t+1 | ρ_t+1 = ρ' ] p(ρ' | ρ  ) dρ' 
       =(y_t = 1 |s_t = s, a_t = a)( (s,a) +∫ V_π( ρ', t+1 ) p(ρ' | ρ∪{s}) dρ' ) 
       + (1 -(y_t = 1 |s_t = s, a_t = a)) ∫ V_π( ρ', t+1 ) p(ρ' | ρ  ) dρ' 
       = ∫ V_π( ρ', t+1 ) p(ρ' | ρ  ) dρ' 
       + (y_t = 1 |s_t = s, a_t = a) ( (s,a) + ∫ V_π( ρ', t+1 ) [ p(ρ'|ρ∪{s}) - p(ρ'| ρ) ] dρ' ) 
       = [ V_π(ρ', t+1) | ρ] 
       + (y_t = 1| s_t = s, a_t = a) ( (s,a) + [ V_π(ρ', t+1) | ρ∪{ s }] - [ V_π(ρ', t+1) | ρ] )

We now use the substitution that (s,a) = C(s)(a - 1). Following the derivation in app:bellman recursion, we define: 

    k_π(s, ρ, t) = 1 - [ V_π(ρ', t) | ρ∪{ s }] - [ V_π(ρ', t) | ρ]/C(s)

Our optimal policy is then the same. Integrating over customers and actions, we obtain the following recursion for the portfolio value function V_π: 

    V_π(ρ, t) = [ V_π(ρ', t+1) | ρ] + [ C(s) (y_t=1|s_t=1,a_t = π(s,ρ,t))(π(s,ρ,t) - k_π(s, ρ, t+1)) ]
 
Given access to p(ρ'|ρ), we can easily adapt alg:training value function to use this recursion instead, by sampling portfolios from p and averaging to approximate the expectations in (<ref>) and (<ref>). This adds some computational overhead, but is not prohibitive. The transition model p(ρ'|ρ) can be learned from historic data. 

One issue we encounter when we account for the possibility of customers leaving the portfolio is that the portfolio objects become considerably more complicated. Recall from sec:RL for PP that, when customers do not leave the portfolio, we can specify portfolios only through their frequency representation, since (i) it was unimportant when a customer joined the portfolio, and (ii) the other customer features were unimportant. However, in this case, neither of this conditions are true. Firstly, the probability of any given customer leaving depends how long they have been in the portfolio. Secondly, the probability of a customer leaving after being in the portfolio for any length of time may be a function of their other features. This necessitates a more detailed and possibly high-dimensional portfolio representation. We do not explore these issues further in our paper. 



 Paper Title: A Vectorization Method Induced By Maximal Margin Classification For Persistent Diagrams.







A Vectorization Method Induced By Maximal Margin Classification For Persistent Diagrams



* Corresponding author

    
An Wu *
College of Mathematics and Computer Science 

Zhejiang A&F University

Hangzhou, China 

20220075@zafu.edu.cn
Yu Pan
School of Mathematics and Statistics 

Beijing Institute of Technology

Beijing, China 

p4nyu@foxmail.com
Fuqi Zhou
School of Electronics and Information 

Hangzhou Dianzi University

Hangzhou, China 

23040444@hdu.edu.cn
Jinghui Yan
Institute of Automation 

Chinese Academy of Sciences

Beijing, China 

jinghui.yan@nlpr.ia.ac.cn
Chuanlu Liu
School of Computer Science & Technology 

Beijing Institute of Technology

Beijing, China 

chuanluliu@163.com

    
==========================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Persistent homology is an effective method for extracting topological information, represented as persistent diagrams, of spatial structure data. Hence it is well-suited for the study of protein structures. Attempts to incorporate Persistent homology in machine learning methods of protein function prediction have resulted in several techniques for vectorizing persistent diagrams. However, current vectorization methods are excessively artificial and cannot ensure the effective utilization of information or the rationality of the methods. To address this problem, we propose a more geometrical vectorization method of persistent diagrams based on maximal margin classification for Banach space, and additionaly propose a framework that utilizes topological data analysis to identify proteins with specific functions. We evaluated our vectorization method using a binary classification task on proteins and compared it with the statistical methods that exhibit the best performance among thirteen commonly used vectorization methods. The experimental results indicate that our approach surpasses the statistical methods in both robustness and precision.



Maximal margin classification, Topological data analysis, Persistent diagram, Vectorization, Protein function




§ INTRODUCTION

Topological data analysis (TDA) is a computational technique which has been developed to study the topological structure of data over the past 20 years<cit.>. The theoretical foundation of TDA is the algebraic topology, and combined with computer technology TDA has been applied to various practical problems, such as image processing, material structure analysis and quantum computing to name a few as in <cit.> and <cit.>.



The primary methods in Topological Data Analysis are persistent homology and Mapper, and we focus on the persistent homology approach in this paper. Based on the persistent homology there are many interesting works. The research team led by Rongling Wu implemented GLMY homology theory proposed by Shing-Tung Yau to analyze and interpret the topological change of health state from symbiosis to dysbiosis, and they successfully described the topological difference and commonality between Crohn's disease and ulcerative colitis<cit.>. Perea introduced how to combine dynamical systems and topological data analysis to gain insights from time-varying data<cit.>. For more applications of persistent homology, one can refer to the review paper of Chazal and Michel<cit.>. Previous works have primarily focused on classification problems as an application of persistent homology. This is because the number of bars in the barcode (Betti numbers) offers a direct measure for classification, or because researchers can define similar functions (such as bottleneck distance and p-Wasserstein distance) based on the persistent diagrams. However, if we want to combine persistent diagrams with more advanced artificial intelligence technologies, there is a critical problem that most machine learning algorithms assume that the input data is from ℝ^n or, more generally, some Hilbert space, but the space of persistent diagrams is not a vector space even if they are from the similar samples. Hence we want to find a  reasonable embedding of persistent diagrams which can be a pre-encoder for any other topological artificial intelligence frameworks.

Many researchers have tried to solve this problem which is called vectorization. Barnes, Polanco and Perea provided a comparative study of six such methods and they applied and compared popular machine learning methods on five data sets<cit.>. Chung and Lawson developed a framework of vectorization and showed several well-known summaries fall in this framework<cit.>. In <cit.> and <cit.>, stable vector representations were developed and these methods were tested on several data sets to show their effectiveness and efficiency. Ali et al. <cit.> summarized thirteen vectorization methods, and they examined these methods on three well-known classification tasks. They discovered that the best-performing method is a simple vectorization by a few statistics. Besides these methods, there are numerous approaches of vectorization as in<cit.> and <cit.>, and more information can be found in Section 2.3 of <cit.>. If we further consider the reasons behind Ali and his colleagues’ results, a plausible explanation is that overly artificial vectorization methods may lead to the loss of key information, whereas statistical methods are relatively more simple and straightforward. However, statistical methods are also hard to interpret. For example, it is difficult to answer why someone need to take the 10th, 25th, 75th and 90th percentiles of the births and deaths rather than others. Hence we approve a more explainable method which is not included in the methods mentioned above, and our method achieve a better result than the statistical method <cit.> which achieves the best performance on several datasets.

As previously stated, one can define a distance function on persistent diagrams to construct a metric space. Consequently, a direct approach is performing classification tasks within this metric space. But a primary challenge mentioned in <cit.> of this method is that 
"The space of persistent diagrams is infinite-dimensional and highly nonlinear",
which poses significant computational and theoretical complexities. However, in practice, since the size of data is always finite, the discrete points in the persistent diagram are bounded. Therefore, we can always construct the  compactification of the space of persistent diagrams generated by the sample set. Finally, by the classical functional analysis theory <cit.>, we have

A compact metric space can always be embedded into a Banach space.

Then we can solve the challenge above on a Banach space.

In the current paper, we focus on the space (𝒳_pd, d) generated by finite persistent diagrams endowed with a fixed metric. Then inspired by the theory of maximal margin classification for Banach spaces <cit.>, we embed (𝒳_pd, d) into the Banach space 𝒞_b(𝒳_pd) by Kuratowski's embedding, then we consider the maximal margin classification for space 𝒞_b(𝒳_pd). As a result, we transform the maximum margin problem into a classical quadratic programming problem as Eq. (<ref>). Moreover, in the derivation process, we find the following result:

Given a set of persistent diagram 𝒳_pd endowed with metric d, we define the vectorization 𝒳_pd→ℝ^n by x↦ (d(x,x_1),⋯, d(x,x_n)), where x_i∈𝒳_pd and n is the cardinal number of 𝒳_pd. In addition, this simple vectorization method corresponds to the quadratic programming problem Eq. (<ref>). 
 

Next, we apply the maximal margin classification for persistent diagrams on the data of Cas-associated proteins, which are closely related to gene editing technology as in<cit.> and <cit.>, and transposases. We choose a simple distance defined by Eq. (<ref>) which requires less computation than the bottleneck distance, and we compare our method with statistical methods which are the best-performing methods in the thirteen vectorization methods <cit.> mentioned before. As a consequent, our method achieves superior performance both in robust and accuracy. 

Finally, we present a framework that uses the TDA processing methodology introduced in this paper to predict proteins with specific functions, and the total framework is exhibited in Fig. <ref>. We will discuss more details in Section III.

The remainder of our paper is organized as follows. In section II, we revisit the maximal margin classification for Banach spaces and apply it to the space of persistent diagrams. Finally, the problem is reduced to a solvable quadratic programming problem. In section III, we show our experiments on the data of Cas-associated proteins and  transposases. We analysis our experiment results and propose a framework of predicting proteins' function. In section IV, we conclusion our work and do some discussions on the current and future researches.



§ MAXIMAL MARGIN CLASSIFICATION METHOD

In this section, we will introduce the main idea of TDA first. Then we will review the principle of the maximal margin classification for Banach spaces, and one can check the details in the publication of Hein, Bousquet and Schölkopf<cit.>. In particular, we will regard the space of persistent diagrams as a Banach space, which allows us to analyze the space of persistent diagrams by maximal margin classification. 



 §.§ The topological data analysis

The main idea of TDA is to find some topological features to describe the difference of data. As an example, assume we have a deflated balloon. When we inflate it, it remains a balloon, because this process is a continuous transformation. However, if we blow in enough gas to cause an explosion, it ceases to be a balloon and becomes a ”burst balloon”. Hence the topological feature that distinguishes burst balloons from intact balloons is the presence or absence of holes in the balloons. By popularizing this idea, researchers have developed a major technique in TDA called persistent homology<cit.>. The word "persistent" means that there is a continuous variation of a parameter which is call the filtration. If we can construct a series of complexes suitable with the continuous variation, then the homological information of the complexes illustrate the topological structure of data. For example, Fig. <ref> shows the process of increasing the radius of points, and it is easy to find that the topological structure of top left three points is different to that of bottom right two points, as the former would create a hole.


Based on this idea, persistent homology is good at dealing with classification tasks, and the topological features are usually taken as the connect components, holes, and voids which are denoted by H_i, i=1,2,3 respectively. Typically, there are two visualization methods for this information, known as the persistent diagram and the barcode, which are essentially equivalent to each other. In Fig. <ref>, we show both visualizations of topological information extracted from Fig. <ref>. 


The construction of complexes in persistent homology is of importance. In this section, we use the Čech complex to illustrate the geometric explanation. In practice, the Vietoris-Rips complex is often employed, as it is beneficial for computational purposes.


However, the number of points contained within different persistent diagrams may indeed differ. Hence, we need a vectorization method to transfer persistent diagrams into a vector space, and we will illustrate our vectorization method next. 




 §.§ The maximal margin classification

Considering the space of persistent diagrams 𝒳_pd, we can define a distance function d(·, ·) on 𝒳_pd to make it a metric space (𝒳_pd,d). In practical problems, the number of generated persistent diagrams is always finite, then we can assume that d(X, Y) is bounded for any sample X, Y∈𝒳_pd, and we use the same symbol 𝒳_pd to denote the subspace generated by sample set. Hence we obtain a compact metric space (𝒳_pd, d), which can be embedded into a Banach space isometrically by the observation of Kuratowski as in<cit.>.

Let 𝒞_b(𝒳_pd) be the Banach space of continuous and bounded functions on 𝒳_pd endowed with infinite norm ._∞. Take an arbitrary point x_0 ∈𝒳_pd and define a map from 𝒳_pd to real-valued functions on 𝒳_pd  

    ϕ : x↦ϕ_x:=d(x, ·)-d(x_0,·).

Then consider the space D :=span{ϕ_x:x∈𝒳_pd}, in which the closure is taken in (𝒞_b(𝒳_pd), ._∞), we have the following lemma: 

  ϕ is a total isometric embedding from (𝒳_pd,d) into the Bananch space (D, ._∞)⊂ (𝒞_b(𝒳_pd), ._∞).


Through Lemma <ref>, we can transform the data processing problem on space (𝒳_pd, d) to that on space (D, ._∞).  In the current paper, we focus on the maximal margin classification of persistent diagrams. Recall that, the margin formulation on a Banach space is

    min_W'∈ D', b∈ℝW'
        s.t.   y_j(⟨ W', ϕ_x_j⟩_D',D+b)≥ 1, ∀ j=1,⋯, n,

where D' is the dual space of D, y_j is the label corresponding to ϕ_x_j and the product ⟨ W', ϕ_x_j⟩_D', D is naturally defined by W'(ϕ_x_j).

Since 𝒳_pd is compact, 𝒞_b(𝒳_pd) is isometrically isomorphic to the Banach space ℳ(𝒳_pd) of finite signed Borel measures on 𝒳_pd with the measure norm. However, we did not construct the parametrization of ℳ(𝒳_pd), hence we need to choose a reasonable approximation on dual space 𝒞_b(𝒳_pd)' instead.

Define a space E:=span{δ_z: z∈𝒳_pd}, where δ_z(f):=f(z) for any f∈𝒞_b(𝒳_pd). There is the following lemma:

  The space E is weak*-dense in 𝒞_b(𝒳_pd)' endowed with the norm 
    ∑_i=1^nβ_iδ_z_i_𝒞_b(𝒳_pd)'=∑_i=1^n|β_i|.


Because of the Lemma <ref>, we can use e=∑_i=1^nβ_iδ_z_i∈ E to approximate any elements of 𝒞_b(𝒳_pd)'. Substitute e into W' in Eq. (<ref>), the optimization problem can be rewritten as 

    min_n∈ℕ, z_i∈𝒳_pd, b∑_i=1^n |β_i|
        s.t.   y_j(∑_i=1^nβ_i(d(x_j, z_i)-d(x_0, z_i))+b)≥ 1.

In consideration of computational convenience, practically it is necessary to restrict E to the training set, which means n is the size of training set. Moreover, it can be noted that the term -∑_i=1^nβ_id(x_0, z_i)+b is independent of j, so that we can reparameterize it as c. In addition, the vector β=(β_1, ⋯, β_n) is optimized in a n dimensional vector space, where dimension n is always finite. Hence the 1-norm in the optimization function is equivalent to 2-norm, which is more advantageous in solving our problem. In light of these observation, we convert the optimization problem as Eq. (<ref>) into solving the following optimization problem:   

    min_β_i, c∑_i=1^n β_i^2
        s.t.   y_j(∑_i=1^nβ_id(x_j, z_i)+c)≥ 1, ∀ j=1,⋯, n.

Observe the form of the optimization problem above, one can find it exhibits a form identical to that of the support vector machine defined in Euclidean space. 

Now we will subsequently explain this phenomenon from the perspective of data science. It is obvious that the primary difference lies in the replacement of term ∑_i=1^nβ_ix_j in SVM with term ∑_i=1^nβ_id(x_j, z_i) in the conditions. Because in SVM the data x_j is always a vector, but now x_j is just a point in the sample space which may no longer be a vector space, therefore we can not do linear combination of samples here directly. However, for each pair of sample points (x_j, z_i) we can define the distance (or similarity) function between them to make the sample space be a metric space. Then for each sample we can generate a fixed-length vector by computing distances from this sample to all points in training set. Essentially, we transform the metric space into a linear space through the distance function, and this transformation method is theoretically guaranteed to be non-arbitrary in mathematics.

Same as SVM, we can also introduce slack variables to obtain a new version with a soft margin as follows:

    min_β_i, ξ_i, c   ∑_i=1^n β_i^2+a∑_i=1^nξ_i
    
    s.t.      y_j(∑_i=1^nβ_id(x_j, z_i)+c)≥ 1-ξ_j, 
       ξ_j≥ 0, ∀ j=1,⋯, n,

where ξ_j is the slack variable which is defined by hinge loss 

    ξ_j=max(0, 1-y_j(∑_i=1^nβ_id(x_j, z_i)+c)).

In the current paper, we focus on the soft margin version of maximal margin classification, and solving Eq. (<ref>) is a typical quadratic programming problem. Here, we render it in matrix notation for ease of subsequent resolution.

Define the coefficient vector β:=(β_1, ⋯, β_n)^T, the slack vector ξ:=(ξ_1,⋯,ξ_n)^T, label vector y:=(y_1, ⋯, y_n)^T and denote the distance matrix of training set by d. Then the optimization function is represented by

    min_β,ξ, c1/2β^T2I_nβ+a 1_n^Tξ,

where I_n is identity matrix with n order and 1_n^T is a n-dimensional vector with elements are all equal to 1. Furthermore, combine β, ξ, and c together, we obtain

    min_α1/2α^TQα+b^Tα,

where α, Q, b are defined as follows:

    α := ([ β; ξ; c ]),  b :=([    0; a1_n;    0 ]), 
        Q := ([ 2I_n    0    0;    0    0    0;    0    0    0 ]).

If we denote the Hadamard product of matrices by the symbol *.
Then similarly, the matrix representation of constraint conditions is

    Gα≥ h,

where the matrix G and the vector h are defined by

    G:=([ y*d I_n 1_n;   0 I_n   0 ]),   
    h:=([ 1_n;   0 ]).


In a conclusion, we transform the problem of maximal margin classification on the space of persistent diagrams into solving the quadratic programming problem

    min_α1/2α^TQα+b^Tα,
        s.t.     Gα≥ h,

which can be solved directly.



§ EXPERIMENTS AND ANALYSIS

In this section, we will illustrate our experiments of maximal margin classification of persistent diagrams. We will apply our Banach space based method, which will be referred as BS method in the following context, to the classification of proteins. Then we will compare the experiment results of BS method with two widely used methods which are based on dimension reduction techniques. Finally, I will analyze the significance of BS method and its applications in the exploration of novel proteins.



 §.§ Experiment Setup

Our experimental setup employs a server equipped with an Intel® Xeon® CPU E7-4870 operating at a frequency of 2.40 GHz, and we use the Python library ripser <cit.> to compute persistent homology. For solving the quadratic programming problem, we use the quadratic programming solver CVXOPT.




 §.§ Dataset

The original dataset we deal with consists of 161 protein PDB files of Cas-associated proteins and 197 files of transposases, and our goal is to classify these proteins correctly. Each PDB file includes the coordinates of atoms in the proteins, and it is easy to extract these coordinate information from PDB files. After extracting the atomic coordinates, the traditional approach of TDA is computing the persistent homology of the atomic point clouds directly. However, this methodology encounters computational challenges, particularly in the calculation of the second persistent homology group (H_2), which is often exceedingly difficult and necessitates substantial computational expenditures. The main reason is that compared to general chemical molecules and structural materials, proteins are typically composed of tens of thousands of atoms. But when calculating homology information of H_2, it is necessary to traverse all combinations of four atomic interactions. Therefore, in this paper, we modify the persistent homology method at the beginning.



 §.§ Preprocessing

We extract the topological information of protein atomic point clouds according to the following approach. Firstly, we derive the coordinates of 358 proteins from the input files. Subsequently, we construct the residue connection matrix of each protein based on the relative positions of the Cα atoms. For more details, we consider the Cα atoms within a protein as nodes, and an edge is created between two Cα atoms if the inter-atomic distance is below this established threshold of 5 angstroms (Å). Following the construction of  residue connection matrices, we can use graph embedding methods to generate point clouds which include less points than the original atomic point clouds. There are many graph embedding methods, and here we choose Node2vec <cit.> method to generate new point clouds. After the embedding of residue connection matrices, we calculate the persistent homology of point clouds including H_0, H_1 and H_2. With the topological features captured, the final step is preprocessing of these topological information. We set a cutoff of 0.01 to filter out H_1 and H_2 information whose persistent times are shorter than this value. We summarize the topological feature extraction algorithm as Algorithm 1.



By using the residue connection matrix and graph embedding, we reduce the number of points in the point clouds. Hence, it is foreseeable that the aforementioned approach to extracting topological features will have a substantial time-saving advantage. In addition, because we only use the inter-atomic distances of the Cα atoms, some information is inevitably omitted, which is acceptable in our experiment.

The results of preprocess is that the computation time required for calculating H_1, H_2, and H_3 using the modified persistent homology approach is comparable to the time taken for calculating H_1 and H_2 using the original persistent homology method, with both processes approximately lasting around 10 days. Roughly speaking, this implies that the new persistent homology algorithm omits certain two-dimensional topological information, and the computational resources saved are redirected towards the calculation of three-dimensional topological information. Although not the focus of this paper, it is worth mentioning that this strategy has a distinct advantage in that it introduces three-dimensional information that was previously unavailable, which will finally provide an additional weight parameter for learning in other machine learning processes.



 §.§ Baselines and Main Method

After the modified persistent homology, we will get persistent diagrams with three components H_1, H_2 and H_3 which are essentially point sets. Next, we will introduce three widely used topological features analysis methods which are based on different statistics as baselines for our BS method as follows:


  * Statistical Method I<cit.>: We count the cardinality of each point set for persistent diagram PD_i, which is equivalent to the persistent Betti number, directly to construct a 3-dimensional feature vector v_i. And the first element is normalized by multiplying with a factor 0.01. 

  * Statistical Method II: 
    We compute the statistical information (the maximum, the minimum, the variance, the mean value and the median) for 5 indicators, corresponding to the birth or death times of H_1, H_2 and H_3, to generate the feature vector of each persistent diagram. Then we extract a 25-dimensional feature vector for each persistent diagram. 

  * Statistical Method III: We apply the feature vectors proposed in the Definition 3.1 of the survey paper <cit.>. This statistical vectorization method incorporates more statistical information than the two methods mentioned above, and it is the best-performing method of the thirteen alternative methods in three well-known classification tasks.


In comparison with the statistical methods, the BS method needs a distance function d to make the space of persistent diagrams 𝒳_pd be a metric space. In this paper, we define the distance function as follows:

    d(X,Y):=∑_i=1^3 w_id_i(X, Y),  ∀ X,Y∈𝒳_pd,

where d_i(X,Y) 

is the maximum Euclidean distance between H_i-components 
of persistent diagrams X and Y, and W=(w_1,w_2,w_3) is weight vector which is fixed as (1/3, 1/3, 1/3) in our experiment.  

In order to compare the effects of different vectorization methods, we apply all methods on the dataset by soft-SVM to obtain classification results.



 §.§ Results and Analysis

We perform five-fold cross-validation experiments and continuously increase the proportion of training data for all methods mentioned before. The experimental results are presented in the following Table I.



The accuracies of the BS method and Statistical methods are noted to increase with the expansion of the training set size. While the Statistical method I has relatively small change around 84%-85%, which means this method exhibits robustness. As a comparison, the Statistical method II and III are more sensitive to the size of training date, and these methods have lower accuracies in the case with small training set. It is worth mentioning that Statistical methods require manual feature selection, and this example demonstrates that different feature selection approaches have a significant impact on the accuracy. Therefore, it needs a prior understanding of the original data which will impact the feature selection. 

However, our BS method uses the information of all features and has a fixed processing procedure, which reduces errors that may arise from poor feature selection. Moreover, the BS method also has robustness, and when the training date size reaches eighty percent, BS method becomes the best method in accuracy. More importantly, the BS method is more generalizable. Because for any persistent diagrams with more topological information (H_4, H_5,⋯), we can also define a distance function like Eq. (<ref>) and we can apply the neural network to determine the weight vector W in practice. 



 §.§ Case Study

In the end, we will illustrate the classification of proteins in the view of biology and propose a framework of finding new proteins with certain functions by topological data analysis. 


For giving a biological explanation, we take 80% of the data to be the training set randomly and use the remaining data to test as well as to build an evolutionary tree. In Fig. <ref>, we show the evolutionary tree and in experiment result we find that 7 proteins have been misclassified. These proteins are 'z7m9a', 'c5vgb', 'c6eyy', 'c6jhv', 'c6jhw', 'c7msl' and 'c7tun', in which we use the characters “c” and “z” to denote Cas-associated proteins and transposases respectively.

It is obvious that, Fig. <ref> can be primarily divided into two components: the upper component mainly consists of Cas-associated proteins, while the lower component mainly consists of transposases. And we find that classification errors can be divided into two parts: 


  * Type I: The first involves the occurrence of another class of proteins within a particular category (such as 'c5vgb', 'c6eyy', 'c6jhv', 'c6hjw' and 'c7msl').

  * Type II: The second involves the presence of errors in the crossover region in which both classes of proteins appear alternatively    (such as 'z7m9a' and 'c7utn'). 


However, if we focus on the instances of misclassification, we find that the reason for Type I error is that the sample classification can be further refined, whereas the reason for Type II error is that the sample classification is overly granular. In details, Cas-associated proteins can be further subdivided into Cas proteins and anti-Cas proteins. Particularly, we find that proteins 'c5vgb', 'c6eyy', 'c6jhv', 'c6jhw' and 'c7msl' are anti-Cas proteins as in<cit.> and <cit.> which are distinct from the majority of other Cas-associated proteins, hence they are misclassified. On the other hand, proteins 'z7m9a' and 'c7utn' are Cas proteins as well as transposases as in <cit.> and <cit.>, and the fact suggests that this dual nature has resulted in their misclassification. Therefore,  our classification methodology aligns with the established principles and demonstrates superior performance.



 §.§ Framework for Protein Function Prediction

Moreover, Type II error approves a new method for discovering novel proteins with specific functionalities. For example, if a transposase in the crossover region is misclassified, we can suspect reasonably that it exhibits structural and functional features similar to those of Cas proteins. Finally, we summarize this framework for exploring protein functions by their spatial structures in Fig. <ref>:





  * In the first step, we input the PDB files {f_j} of proteins with a known specific function and {g_i} of validation proteins;

  * In the second step, we calculate the topological information for every PDB file from its spatial structure;

  * In the third step, we conduct the binary classification task described in Section II to differentiate between the sets {f_j} and {g_i} using our vectorization method of persistent diagrams;

  * Finally, we confirm those proteins with the second type of error and identify the subset {g_k} of proteins that exhibit this function.



 



§ CONCLUSION AND DISCUSSION



 §.§ Conclusion

We propose a novel vectorization method, which not only attains the best performance in processing topological information compared with conventional methods but also exhibits robustness against the size of the training dataset, to analysis persistent diagrams. We derive the process for handling persistent diagrams using the maximal margin classification method of a Banach space, and discover that once the metric space of persistent diagrams is embedded by Kuratowski's embedding then the final result is equivalent to an encoder that calculates the distances between a given persistent diagram and all other diagrams. Hence our result proposes a more reasonable feature selection which use more topological information extracted than statistical methods that are the best method between thirteen alternative
methods<cit.>. 

Another result with more profound implications for biology is that we propose a framework 
of using spatial topological information of proteins to identify their functions. Because we find that in our binary classification experiment, the second type of error typically arises from proteins that belong to both classes. Therefore, to verify whether certain proteins possess a specific function, we need only focus on whether they would result in a Type II error when classified alongside proteins which are known to exhibit this function.

As a byproduct of our processing, we have optimized the persistent homology algorithm for protein analysis, therefore its computational efficiency is enhanced and it is permitted to calculate the second Betti number.



 §.§ Discussion

In future work, we plan to investigate that whether the Banach space based method can be improved by following modifications:

  
  * First, we intend to examine the performance of our method across various distances;
  
  * Second, we plan to combine our method with neural network methods;
  
  * Finally, we aim to propose additional vectorization methods for persistent diagrams, informed by geometric perspectives


In the derivation process of the BS method, we used information from the distance function. 
Compared to traditional methods of comparing persistence diagrams, such as the bottleneck distance, 
the theoretical computational complexity of our method is lower 
because it does not require enumerating all possible matchings. 
However, we may get a better performance distance by training the metric matrix with the neural network.


Our method can be integrated with the neural network in more aspects, because our method maps proteins into a vector space of relative topological information, 
and this vector representation 
offers greater interpretability. In our subsequent work, we will combine the vectorization approach presented in this paper with neural networks to address further problems related to the spatial structure of proteins. 








Finally, since the vectorization method proposed in this paper originates from a geometric embedding, we aim to consider more vectorization methods inspired by other geometric structures such as Lie groups and homogeneous spaces. 
Particularly, we can choose different embedding methods than Kuratowski's embedding to construct different reasonable encoders for persistent diagrams. 



§ ACKNOWLEDGEMENTS

The authors  extend many thanks to Professor Chao Qian and Hanning Wang who provide academical support and encouragement. 

1
AAJ23
D. Ali, A. Asaad, M. J. Jimenez, V. Nanda, E. Paluzo-Hidalgo, and M. Soriano-Trigueros, “A survey of vectorization methods in topological data analysis," IEEE Trans. Pattern Anal. Mach. Intell., vol. 45, no. 12, pp. 14069–14080, 2023.

BPP21
D. Barnes, L. Polanco, and J. A. Perea, “A Comparative Study of Machine Learning Methods for Persistence Diagrams," Front. Artif. Intell., vol. 4, pp. 681174, 2021.

B21
U. Bauer, “Ripser: efficient computation of Vietoris–Rips persistence barcodes," J Appl. and Comput. Topology, vol. 5, pp. 391–423, 2021.

C09
G. Carlsson, “Topology and data," Bull. Amer. Math. Soc., vol. 46, no. 2, pp. 255-308, 2009.

C20
G. Carlsson, “Topological methods for data modelling," Nat. Rev. Phys., vol. 2, pp. 697-708, 2020.

CM21
F. Chazal and B. Michel, “An Introduction to Topological Data Analysis: Fundamental and Practical Aspects for Data Scientists," Front. Artif. Intell., vol. 4, pp. 667963, 2021.

CLW23
D. Chen, J. Liu, J. Wu, G.-W. W, F. Pan, and S.-T. Yau, “Path Topology in Molecular and Materials Sciences," J. Phys. Chem. Lett., vol. 14, no. 4, pp. 954-964, 2023.

CNO20
I. Chevyrev, V. Nanda, and H. Oberhauser, “Persistence paths and signature features in topological data analysis," IEEE Trans. Pattern Anal. Mach. Intell.,  vol. 42, no. 1, pp. 192–202, 2020.

CL22
Y.-M. Chung and A. Lawson, “Persistence Curves: A canonical framework for summarizing persistence diagrams," Adv. Comput. Math., vol. 48, no. 6, pp. 1-42, 2022.

CRC13
L. Cong, F. A. Ran, D. Cox, S.-L. Lin, R. Barretto, N. Habib et al., “Multiplex Genome Engineering Using CRISPR/Cas Systems," Science, vol. 339, no. 6121, pp. 819-823, 2013.

DHZ20
Z.-T. Dong, C.-F. Hu, C. Zhou, and H.-W. Lin, “Vectorization of persistence barcode with applications in pattern classification of porous structures," Comput. Graph., vol. 90, pp. 182-192, 2020.

DLZ24
Z.-T. Dong, H.-W. Lin, C. Zhou, B. Zhang, and G.-C. Li, “Persistence B‑spline grids: stable vector representation
of persistence diagrams based on data fitting," Machine Learning, vol. 113, pp. 1373-1420, 2024.

ELZ02
H. Edelsbrunner, D. Letscher, and A. Zomorodian, “Topological Persistence and Simplification," Discrete Comput. Geom., vol. 28, pp. 511–533, 2002.

GL16
A. Grover and J. Leskovec, “node2vec: Scalable Feature Learning for Networks," KDD '16: Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, 2016.

HDM17
L. B. Harrington, K. W. Doxzen, E. Ma, J.-J. Liu, G. J. Knott, A. Edraki et al., “A Broad-Spectrum Inhibitor of CRISPR-Cas9," Cell, vol. 170, no. 6, pp. 1224-1233.e15, 2017.

HBS05
M. Hein, O. Bousquet, and B. Schölkopf, “Maximal margin classification for metric spaces," J. Comput. Syst. Sci., vol. 71, no. 3, pp. 333-359, 2005.

HNH16
Y. Hiraoka, T. Nakamura, A. Hirata, E. G. Escolar, K. Matsue, and Y. Nishiura, “Hierarchical structures of amorphous solids characterized by persistent homology," Proc. Natl. Acad. Sci. USA., vol. 113, no. 26, pp. 7035-40, 2016.

HPG22
S. Hwang, C.-X. Pan, B. Garcia, A. R. Davidson, T. F. Moraes, and K. L. Maxwell, “Structural and Mechanistic Insight into CRISPR-Cas9 Inhibition by Anti-CRISPR Protein AcrIIC4Hpa," J. Mol. Biol., vol. 434, no. 5, pp. 167420, 2022.

HRA18
A. P. Hynes, G. M. Rousseau, D. Agudelo, A. Goulet, B. Amigues, J. Loehr et al., “Widespread anti-CRISPR proteins in virulent bacteriophages inhibit a range of Cas9 proteins," Nat. Commun., vol. 9, pp. 2919, 2018.

JCF12
M. Jinek, K. Chylinski, I. Fonfara, M. Hauer, J. A. Doudna, and E. Charpentier, “A Programmable Dual-RNA–Guided DNA Endonuclease in Adaptive Bacterial Immunity," Science, vol. 337, no. 6096, pp. 816-821, 2012.

JJ19
P. Joharinad and J. Jost, “Topology and curvature of metric spaces," Adv. Math., vol. 356, pp. 106813, 2019.

K18
S. Kališnik, “Tropical coordinates on the space of persistence Barcodes," Found. Comput. Math., vol. 19, no. 1, pp. 101–129, 2018.

KLY19
Y. Kim, S. J. Lee, H. J. Yoon, N.-K. Kim, B.-J Lee, and J.-Y. Suh, “Anti-CRISPR AcrIIC3 discriminates between Cas9 orthologs via targeting the variable surface of the HNH nuclease domain," FEBS. J., vol. 286, no. 2, pp. 4661-4674, 2019.

LGZ16
S. Lloyd, S. Garnerone, and P. Zanardi, “Quantum algorithms for topological and geometric analysis of data," Nat. Commun., vol. 7, pp. 10138, 2016.

PTM21
J. Park, A. W. Tsai, E. Mehrotra, M. T. Petassi, S.-C. Hsien, A. Ke et al., “Structural basis for target site selection in RNA-guided DNA transposition systems," Science, vol. 373, pp. 768-774, 2021.

PC14
J. A. Perea and G. Carlsson, “A Klein bottle based dictionary for texture representation," Int. J. Comput. Vis., vol. 107, no. 1, pp. 75-97, 2014.

P19
J. A. Perea, “Topological time series analysis," Notice of AMS., vol. 66, no. 5, pp. 686-694, 2019.

SHK22
G. Schuler, C. Hu, and A. Ke, “Structural basis for RNA-guided DNA cleavage by IscB-ωRNA and mechanistic comparison with Cas9," Science, vol. 376, pp. 1476-1481, 2022.

TSB18
C. Tralie, N. Saul, and R. Bar-On, “Ripser.py: A Lean Persistent Homology Library for Python," J. Open Source Softw., vol. 3, pp. 925, 2018.

WLD23
S. Wu, X. Liu, A. Dong, C. Gragnoli, C. Griffin, J. Wu et al. “The metabolomic physics of complex diseases," Proc. Natl. Acad. Sci. USA., vol. 120, no. 42, pp. e2308496120, 2023.

ZKN24
A. Zia, A. Khamis, J. Nichols, U. B. Tayab, Z. Hayder, V. Rolland et al., “Topological deep learning: a review of an emerging paradigm," Artif. Intell. Rev., vol. 57, no. 4, 2024.

ZLJ20
B. Zieliński, M. Lipiński, M. Juda, M. Zeppelzauer, and P. Dłotko, “Persistence codebooks for topological
data analysis," Artif. Intell. Rev., vol. 54, no. 3, pp. 1969–2009, 2020.










 Paper Title: DeepBaR: Fault Backdoor Attack on Deep Neural Network Layers.









: Fault backdoor attack on deep neural network layers
    Camilo A. Martínez-Mejía
Independent Researcher 

Bogotá, Colombia

E-mail: camilo.csmath@gmail.com
Jesus Solano
ETH Zürich

Zürich, Switzerland

jesus.solano@inf.ethz.ch
Jakub Breier
TTControl GmbH

Vienna, Austria

E-mail: jbreier@jbreier.com
Dominik Bucko
Deutsche Telekom Cloud Services

Bratislava, Slovakia

bucko.dominik@protonmail.com
Xiaolu Hou
Slovak University of Technology

Bratislava, Slovakia

houxiaolu.email@gmail.com
    30 July 2024
======================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Machine Learning using neural networks has received prominent attention recently because of its success in solving a wide variety of computational tasks, in particular in the field of computer vision. However, several works have drawn attention to potential security risks involved with the training and implementation of such networks. In this work, we introduce , a novel approach that implants backdoors on neural networks by faulting their behavior at training, especially during fine-tuning. Our technique aims to generate adversarial samples by optimizing a custom loss function that mimics the implanted backdoors while adding an almost non-visible trigger in the image. We attack three popular convolutional neural network architectures and show that  attacks have a success rate of up to 98.30%. Furthermore,  does not significantly affect the accuracy of the attacked networks after deployment when non-malicious inputs are given. Remarkably,  allows attackers to choose an input that looks similar to a given class, from a human perspective, but that will be classified as belonging to an arbitrary target class. 







§ INTRODUCTION


In recent years, machine learning using neural networks has seen significant progress <cit.>. For instance, classification tasks in computer vision have reached high levels of accuracy, and are fundamental for autonomous driving among others <cit.>. Naturally, concerns about the security and safety of such solutions have been raised and studied <cit.>. A natural concern arising in this context is: Can inputs to computer vision classifiers (images), which would be classified correctly by a human observer be crafted by adversaries such that they will be misclassified by machine learning algorithms? A motivation for an attacker to perform such attacks would be to maliciously cause an autonomous vehicle to misbehave or to bypass an authentication mechanism based on face recognition among others <cit.>. Such attacks have been shown to be possible using adversarial learning <cit.>, a technique that exploits the fact that boundaries between classes in a machine learning classifier are easy to cross by small perturbations to inputs. 

However, a much harder problem is to achieve targeted adversarial attacks, where an adversary produces inputs that look (to humans) as if they would belong to a given class but will be misclassified as an arbitrary target class <cit.>. 
The difficulty of targeted attacks comes from the need to navigate the model’s decision boundaries in high-dimensional space to find a specific adversarial example that leads to a specific incorrect classification. This requires a more sophisticated understanding of the model’s internal workings and often requires more computational resources compared to non-targeted attacks where the goal is to find any point that crosses the model’s decision boundary.
In previous work <cit.>, it has been shown possible to achieve a higher degree of success rate by introducing backdoors to neural networks under training by using fault injection attacks <cit.>. Although promising, this approach had several drawbacks. First, its success depended on finding the solution to a constraint-solving problem, which was shown to work only in simple network architectures where fault injections were performed at superficial layers. This technique for instance cannot be directly applied to modern convolutional neural networks such as VGG-19 <cit.>.

In this work, we show a novel technique to insert backdoors in complex convolutional network architectures while achieving a high attack success rate, called . At the core of our technique is a faulting attack similar to the one presented in <cit.> but which does not rely on constraint solving to find the fooling images. Instead, we propose a strategy that relies on generating adversarial samples by optimizing custom loss functions that mimic the fault attacks performed during training while keeping the image similarity, so attacks (i.e. changes in the image) are hardly perceptible by humans. We evaluate our approach on highly popular convolutional neural networks such as VGG-19 <cit.>, ResNet-50 <cit.>, and DenseNet-121 <cit.>. As a result, we obtain high attack success rates, of up to 98.30% for VGG-19, 97.94% for ResNet-50, and 88.94% for DenseNet-121. 

When comparing our approach to traditional targeted adversarial examples <cit.>, it provides several advantages.
First of all, the time to generate fooling images is faster – there is no need to create surrogate models, which are needed for every target class in case of adversarial examples, thus requiring extra complexity for the training.
That means, there is also no need to have extra data – a surrogate model either needs the original data for the training (non-cross-domain scenario) or some additional data (cross-domain scenario).
Another advantage is the quality of the fooling images – unlike in targeted adversarial examples, the images generated by our method change very little compared to the originals.
Finally, we also showed that by performing traditional adversarial training, the  of 's adversarial samples drops significantly for all the datasets and architectures studied. 






In sum, our contributions are:


    
  * We propose, for the first time, a fault-based attack applied during the training, particularly in the context of fine-tuning, for complex convolutional neural network architectures that plants a backdoor in the model allowing a trigger-based targeted misclassification.
    
  * Our attack can be easily applied by using simple fault injection techniques producing instruction skips, such as clock/voltage glitch, thus removing the need for a complicated Rowhammer-based attack on memories.
    
  * Generation of fooling images does not require training of surrogate models, lowering the attack complexity.
    
  * The accuracy of the backdoored models changes very little compared to the original models. Particularly,  reduces the performance of the downstream model by up to 0.8% on average.
    
  * We evaluate our approach and show it has higher attack success rates than the current state-of-the-art weight-oriented approach, DeepVenom <cit.>.
    
  * We propose adversarial training as a countermeasure to our approach showing that after re-training with fooled images generated with  the  drops to roughly 5% regardless the architecture.


The rest of the paper is organized as follows. Section <ref> recaps fundamental machine learning and fault attack concepts. Section <ref> describes our approach and faulting strategy. Section <ref> describes the experimental setup, datasets considered for the evaluation, and the evaluation protocol. Section <ref> presents the results of the evaluation for in-domain and out-of-domain data. Section <ref> shows a set of strategies to defend against . Finally, Section <ref> shows a comprehensive summary of the state-of-the-art approaches to adversarial and fault-based backdoor attacks.






§ BACKGROUND AND PRELIMINARIES



In this section, we will first give an introduction to neural networks (Section <ref>) and fault attacks on neural networks (Section <ref>).
Section <ref> provides an overview of the Structural Similarity Index (SSIM). 
Then in Section <ref> we present the terminologies that will be used for the rest of this paper.
Finally, Section <ref> formalizes fooling backdoor attacks.



 §.§ Artificial neural networks



Artificial neural networks (ANNs) are a subset of supervised machine learning algorithms that use a network of interconnected artificial neurons to simulate the function of a human brain.
Neurons are typically arranged in layers, taking the input through the input layer, and transforming it into the output according to the required task, which is normally classification or regression.
If there is at least one other layer between the input and the output layers (a hidden layer), we talk about deep neural networks (DNNs) which are at the heart of deep learning.
Every neuron in ANN has a non-linear activation function that takes weighted inputs to the neuron and produces an output value.
In our work, we are interested in networks based on ReLU (rectifier linear unit) activation function <cit.>.
ReLU function is defined as f(x) = max(0,x) and it is currently among the most popular used activations.

ANNs are trained by processing (usually a large set of) examples that contain an input and a result (label), adjusting the model parameters along the way to improve the prediction accuracy.
The training algorithm that performs this step is called backpropagation <cit.>.

Convolutional neural networks

Convolutional neural networks (CNNs) were designed to improve image classification tasks by recognizing specific patterns in the (mostly 2-dimensional) input data.
CNNs are ANNs that utilize a form of feature extraction within their convolutional layers <cit.>.
Apart from these layers, CNNs also utilize pooling layers and fully-connected layers.
Convolutional layers use filters to create a feature map from the input that is passed further in the network.
Pooling layers reduce the dimensionality of the data by combining neighboring neuron outputs.
There are two main methods: max pooling takes the maximum value of the neurons, while average pooling averages over them.
Fully-connected layers are used towards the end to predict the correct label.




 §.§ Fault attacks on neural networks



Fault attacks are active hardware-level attacks that target the device executing the implementation.
Originally, they were proposed against cryptographic algorithms, allowing very efficient key recovery attacks <cit.>.
They can either be performed with physical access to the device, for example by means of clock/voltage glitching, electromagnetic pulses, lasers <cit.>, or remotely by using techniques such as Rowhammer <cit.>, CLKSCREW <cit.> or VoltJockey <cit.>.
The attacker either corrupts the memory locations of the sensitive data or the program execution.
In the first case, corruption can lead to bit-flips, bit sets/resets, or random byte faults <cit.>.
In the second case, the attacker can alter the program execution by skipping <cit.> or changing <cit.> some instructions.

The first fault attack proposal targeting neural network models was published in 2017 <cit.>, utilizing bit-flips to cause misclassification.
The first practical attack followed a year after and it used instruction skips caused by a laser to change the output of executed activation functions <cit.>.
After that, the main body of work used Rowhammer faults to flip bits in DRAM memories <cit.>.
The Rowhammer-based methods were termed adversarial weight attacks, and apart from <cit.>, they all targeted quantized neural networks – models that use low-precision data representation, typically with a word-length of 8 or 4 bits instead of a traditional floating point representation (IEEE 754).
Apart from misclassification, other attack vectors have been investigated, such as model extraction <cit.>, trojan insertion <cit.>, and backdoor injection <cit.>.
In this work, we aim at backdoor injection as well, but unlike in <cit.>, our method can scale up to large real-world models, thus bringing this attack vector to the practical realm.


ReLU-skip attack

At the core of this work, there is a fault model that was shown to be practical in <cit.>.
The main idea is to skip a jump instruction during the ReLU execution to always force the function output to be zero, as illustrated in Figure <ref>.
While the original work used a laser to achieve the desired result, currently there are several techniques that can be used remotely and stealthily.
For example, CLKSCREW <cit.> and VoltJockey <cit.> use energy management features in the software to exploit frequency and voltage hardware regulators to remotely skip instructions on ARM processors.
In <cit.>, the authors show how voltage drop from an FPGA located on the same die as a CPU can cause instruction skips.
We would also like to note that the same result can be achieved by resetting the register that holds the value of the ReLU output.
This extends the attack to models implemented on FPGAs, as the register reset on these devices was successfully demonstrated in <cit.>.
In the following, we refer to this fault model as ReLU-skip attack.





 §.§ Faulting assumptions



Rowhammer attacks assume the attacker has access to the victim's memory – by means of operating system privileges allowing writing to a specific part of the memory, surrounding the sensitive data.

Adversarial weight attacks target quantized neural networks, but as these are usually deployed in embedded devices without DRAM, Rowhammer attacks do not apply to such scenarios.

Apart from that, there is a wide body of work dedicated to Rowhammer defenses <cit.>.

On the other hand, causing instruction skips is relatively easy <cit.>. In the context of , we assume that the attacker can execute ReLU-skip attacks, and has access to weights and implementation code of the infected model. This access allows them to leverage its methods and attributes.



 §.§ SSIM



The Structural Similarity Index (SSIM) <cit.> is a metric used to quantify the similarity between two images. Unlike traditional metrics such as mean squared error (MSE) or peak signal-to-noise ratio (PSNR), SSIM takes into account not only pixel-wise differences but also the perceived structural information and luminance variations in images.

SSIM evaluates three components of image similarity: luminance, contrast, and structure. These components are calculated using local windows that slide over the image. Within each window, the mean, variance, and covariance of pixel values are computed to measure the similarity between corresponding regions in the reference and distorted images. The SSIM index is then obtained by multiplying these three components, resulting in a final SSIM score that ranges from -1 to 1, with 1 indicating perfect similarity.



 §.§ Terminology



Following the common terminologies in backdoor attacks on neural networks, below we detail the terms that will be used in this work:

    
  * Benign model is the model trained without fault attacks.
    
  * Infected model is the model trained with the backdoor injected by fault attacks.
    
  * Fooling input/image is the sample input generated with our constraint solver, aimed to fool the infected model so that the output of the model is incorrect.
    
  * Base pattern is the sample used to generate fooling input from. It can be a random sample outside the problem domain.
    
  * Source class is the output of the benign model given an input.
    
  * Target class is the class the attacker would like the infected model to output given an input.
    
    
  * Benign accuracy is the test accuracy of the benign model.




 §.§ Fooling backdoor attack


Let c_t denote our target class of the attack.
Faults are injected during the training of a benign model ℳ and result in an infected model ℳ'.
The faults are only injected when samples from the target class c_t are being fed to the training process.

Given any base pattern 𝐱, with the knowledge of fault locations and effects injected during the training, the attacker constructs a fooling input 𝐱' from 𝐱.
The base pattern 𝐱 can be a legitimate sample from the test set, and can also be an out-of-the-domain sample that belongs to one of the possible classes.

An attack on 𝐱 is said to be successful if 

    ℳ'(𝐱')=c_t.


However, if the classification result of 𝐱' by ℳ' has very low confidence, denoted CF(𝐱'), we would expect the victim to be suspicious about the output.
Hence, we also define a threshold τ that the confidence should reach for a successful attack.
And we say an attack on 𝐱 is successful above the threshold if the misclassification from Equation <ref> has confidence CF(𝐱')>τ.
Similarly, we say an attack on 𝐱 is successful below the threshold if the misclassification from Equation <ref> has confidence CF(𝐱')≤τ.

Let c_𝐱 denote the source class of 𝐱, i.e.

    ℳ(𝐱)=c_𝐱.

If Equation <ref> holds and we also have

    ℳ(𝐱')=c_𝐱,

then we say that the successful attack on 𝐱 has successful validation.

Let S denote a set of N base patterns

    S={𝐱_1,𝐱_2,…,𝐱_N}.

And let S' be the set of corresponding fooling inputs

    S'={𝐱'_1,𝐱'_2,…,𝐱'_N}

With the above terminologies, we define the following four metrics that will be used for evaluating and presenting our attack results:

    
  * Attack Success Rate Above Threshold, denoted , is given by the percentage of inputs 𝐱 in S such that the attack on 𝐱 is successful above the threshold.
    In other words, 
    
    =|{i | ℳ'(𝐱_i')=c_t,CF(𝐱')>τ, i=1,2,…, N}|/N

    
  * Attack Success Rate Below Threshold, denoted , is given by the percentage of inputs 𝐱 in S such that the attack on 𝐱 is successful below the threshold.
    Or equivalently, 
    
    =|{i | ℳ'(𝐱_i')=c_t,CF(𝐱')≤τ, i=1,2,…,N}|/N

    
  * Attack Success Rate, denoted , is given by the percentage of 𝐱 in S such that the attack on 𝐱 is successful.
    Namely,
    
    =+







§ OUR APPROACH



Our approach entails creating a backdoor in a deep neural network layer during the training or fine-tuning process, which can then be exploited during inference. We have two phases: the faulting strategy (Section <ref>) and the fooling image generation strategy (Section <ref>). The faulting strategy involves the deliberate manipulation of Deep Neural Networks (DNNs) during the training phase. Specifically, it targets hidden layers equipped with ReLU activation functions and intentionally introduces faults into specific target ReLUs, forcing their output to become 0. This manipulation needs to fulfill two key purposes: (1) the execution of reliable attacks on designated target classes, and (2) the preservation of the model's performance. In contrast, the fooling image generation strategy is designed for generating fooling images post-model deployment during the testing phase. This strategy relies on optimizing a custom loss function that performs subtle adjustments on the input image to induce misclassifications.







 §.§ Faulting strategy



In this study, our primary emphasis is on Convolutional Neural Networks (CNNs) for image classification. Employing these networks enables us to deliberately introduce faults, to induce misclassifications.

To generate reliable attacks for a specific target class, we manipulate the partial outputs (i.e., the computations that follow a layer) during the training phase. This manipulation involves a physical action: executing the ReLU-skip attack detailed in Section <ref>. Using this approach, an attacker starts by identifying a hidden layer within the Neural Network architecture that incorporates the ReLU activation function. Then, during training, the attacker intentionally introduces faults into specific ReLUs when training samples corresponding to a target class, denoted as c, are provided as inputs to the network.
Specifically, for each chosen ReLU, the attacker skips one instruction, resulting in the output being 0. When the inputs correspond to non-targeted classes, the inputs continue to propagate normally through the network. In other words, we do not execute attacks in such cases.

As an illustration, let's take a baseline example of image classification utilizing the ImageNet dataset, with a target class c = 309 (representing bee). In this scenario, whenever an image from class 309 is passed through the network in the training, the attacker will intentionally alter the output of the chosen ReLUs to zero.

Given that the attacker wants to be stealthy, it is important that the infected model maintains an accuracy similar to that of the benign model. The attacker's strategy will be to fault all the ReLUs of the targeted layer, while at the same time selecting a fraction, represented by χ, of the training samples from the chosen target class to be subjected to faults. Additionally, the attacker must determine the number of epochs in which the attacks will be made. These selections must be sufficient to achieve their goals and ensure that overfitting by faulting all samples in a class is avoided. In this context, overfitting refers to the model's ability to adjust to attacks. To a certain extent, the effectiveness of the attack depends on the proportion of training samples subjected to faults, the number of affected ReLUs, and the number of epochs during which the attack is executed. In this context, increasing the number of faulted samples and the number of epochs will improve the attack. Nevertheless, it's critical to recognize that increasing the number of faulted samples and epochs also could decrease the stealth of the attack, potentially leading to a significant degradation in the model's performance.








 §.§ Fooling images generation strategy



Notice that the faulting strategy targets the training phase. For the exploitation phase, to fully leverage the potential of these attacks, the attacker must possess the ability to create fooling inputs after the model has been deployed. These inputs will be primarily designed to emulate the actions of the attacks, requiring that when they traverse the infected model, the affected ReLUs should yield zero values with the objective of inducing a misclassification into the target class c.

To generate the fooling inputs, following the assumptions outlined in Section <ref> for , we propose that this process can be achieved by optimizing a custom loss function tailored for images. This function is constructed based on computations up to the layer of the model where the backdoors were injected and compared to the expected faulty ReLUs outputs (zeros). These computations within the model are performed on a specific fragment, which we denote as model ℳ'|_1^λ or, for simplicity, ℳ_λ', since the attacker always need to feed the network from the first layer, so the subscript is not needed. Here, λ represents the index of the compromised layer with the associated ReLU.

The optimization process is carried out with respect to the input of ℳ_λ'. The input of the model ℳ_λ' is 𝐱', which is initialized as a copy of 𝐱, so 𝐱' will change through the iterations of the optimization model until the confidence level exceeds a threshold τ or until a maximum number of iterations μ is reached. The threshold τ is the value that the confidence level should reach for a successful attack. Our custom cost function, which we will denote as ℒ_fool, consists of two complementary functions. The first, designated as ℒ_out, corresponds to the Huber Loss with the target value set to zero and the reduction operation set to the sum. This loss function is defined as 

    ℒ _out =
            
                0.5 · (x_n - y_n)^2,    if   |x_n - y_n| < δ
    δ· (|x_n - y_n| - 0.5 ·δ),    otherwise.

In particular, we empirically set δ to 0.5, and due to the attack structure, the target is 0, thus y_n=0. ℒ _out allows input adjustments to ensure that the compromised ReLUs produce zeros as intended.

On the other hand, the second function, ℒ _ssim, relates to the structure similarity index SSIM, which measures the similarity between two images. Therefore, ℒ _ssim  is outlined as  


    ℒ _ssim = 1 - SSIM(𝐱',𝐱).

ℒ _ssim aims to keep the manipulated inputs as close to the base pattern 𝐱 as possible. 

Therefore, our custom loss function for generating fooling images is the sum of ℒ _out and ℒ _ssim, so 

    ℒ_fool = ℒ _out + ℒ _ssim.

This provides a means to control the amount of change applied to the input 𝐱'. Initially, ℒ _out and ℒ _ssim are on different scales. ℒ _ssim starts at 0, since 𝐱 and 𝐱' are identical, and then remains very close to 0. Conversely, ℒ _out represents a type of error relative to the target 0. However, as optimization progresses, ℒ _out converges to a similar scale. This characteristic allows ℒ to primarily focus on exploiting the injected vulnerability.

In summary, ℒ_fool serves a dual purpose: it enables the exploitation of the injected vulnerability while ensuring that 𝐱' closely resembles 𝐱 in human perception.




§ EVALUATION



Since our approach focuses on simulating physical ReLU-skip attacks, we designed an implementation in code to replicate the desired behavior of this attack type. Subsequently, we apply the fooling image generation strategy described in <Ref>. Finally, to evaluate the effectiveness of our approach, we report two of the attack metrics mentioned in <Ref>, , and . Following the evaluation protocol of previous work in the field <cit.>, the results are presented as the average of the  of three target classes from the ImageNet dataset: Great Grey Owl, Goose, and French Bulldog. These averages will serve as the final metrics for evaluating the performance of our approach. It should be noted that some metrics depend on a confidence threshold τ. For our evaluation, we empirically set τ to 0.1. A confidence value of 0.1 for the threshold τ means that the model assigns more than 10% of probability to the predicted class. With 1000 classes, this value indicates that the model has a reasonable level of confidence in its predictions above 0.001 when all classes have the same probability.

The efficacy of  is evaluated, as in <cit.>, on three widely-used network architectures, including VGG-19 <cit.>, ResNet-50 <cit.>, and DenseNet-121 <cit.>.



 §.§ Experimental Setup



In this work, we extend the training of models previously trained on the ImageNet <cit.> dataset, which can be considered a form of fine-tuning with the same dataset. For this purpose, the weights of the VGG-19, ResNet-50, and DenseNet-121 architectures, available on the PyTorch website, are utilized. 

Datasets: Building upon previous work <cit.>, we utilize two datasets in our experiments: the ImageNet <cit.> and the Paintings <cit.> datasets. The provided training and validation sets of ImageNet are employed. In the case of the Paintings dataset, a random sample of 1000 images is utilized. 

Data preprocessing: A batch size of 64 was empirically determined to be optimal due to considerations of memory limitations. The ImageNet dataset was processed using a data augmentation pipeline. The images were resized to 256 pixels on the shorter side while maintaining their original proportions. Additionally, random horizontal flipping and center cropping were applied to achieve a final size of 224×224 pixels. Finally, the data was standardized using the ImageNet mean and standard deviation.

Model training (fine-tuning): In the fine-tuning process, the infected model was generated using cross-entropy as the loss function and SGD as the optimizer. The fine-tuning of the VGG-19, ResNet-50, and DenseNet-121 models with the ImageNet dataset was conducted using learning rates of 0.0001, 0.001, and 0.001, respectively, with a momentum of 0.9 and a weight decay of 5×10^-4. The cosine annealing schedule was incorporated as the learning rate. Two major experiments were conducted. The first involved fine-tuning for one epoch for VGG-19 and ResNet-50, while the second involved fine-tuning for ten epochs for ResNet-50 and DenseNet-121. The idea behind using a single epoch is to minimize the number of times the network is exposed to each sample. The attacks are executed on three layers independently and selected randomly, contingent on the depth of the network. Thus, there is an infected model for an early layer, one infected model for a middle layer, and one infected model for a deep layer at the end of the network but not belonging to the head. 





Fooling image generation: At this point, we have the infected models. Then we need to exploit the vulnerability. Subsequently, the vulnerability must be exploited. To achieve this, the objective function, which needed to be minimized, was defined by Equation <ref> and explained in Section <ref>. Adam was employed as the optimizer, with a learning rate of 0.015 for VGG-19 and 0.025 for the other networks. These learning rates were chosen experimentally. The maximum number of iterations μ was empirically determined to be 200. In this case, the model is only able to see the sample 200 times (queries).

Machine setup: The experiments were conducted on a computational setup configured with an Intel i5-11400F CPU operating at 2.6 GHz. The system was equipped with 16 GB of DDR4 RAM clocked at 3200 MHz. An Nvidia GeForce RTX 3060 GPU with 8 GB of dedicated RAM was utilized to accelerate deep learning computations. 



 §.§ Attack overview



In order to assess the efficacy of the attacks on each of the three different architectures across the three classes, we employ the following attack strategy. First, a model is selected. Second, a target class is selected for the faulting process. Third, the layer to be attacked is chosen. Fourth, the attack is executed. In this instance,  we set the value of χ to 0.9. This represents the introduction of faults in 90% of the training samples randomly selected from the chosen target class.

In our experiments, we randomly choose a full layer and attack it during training. Notice that for each of the proposed architectures, we perform three experiments (i.e., we perform the attack on three different layers independently). For DenseNet-121, only the final ReLU is attacked. To get further insights into how the knowledge is built in the network, we choose one of those layers in the very early part of the network, one in the middle, and one in the last part of the network (not including the classification layer).

Once the process of fine-tuning while faulting the network is completed, the infected model is ready for the fooling image generation, as detailed in Section <ref>. In this phase, only a single image is required. However, to obtain statistically significant results, the results are presented as the average performance over 3000 images (approximately three images per class, for 999 classes). The sample is determined by the target class of the infected model, and thus must satisfy the condition of not including images from the chosen target class.







§ RESULTS






 §.§ ImageNet Dataset



The results are presented in the two phases of : the faulting attack and the image generation. In the <Ref> we present the results for the three evaluated architectures on the data domain of the models (ImageNet). In this scenario, attackers and users share the same data domain. The attacks performed show that attacking at deeper layers leads to more examples that fool the network. Remarkably, we can see that  achieves a  superior to 98% after performing the attack for only one epoch, especially in the case of VGG-19. In contrast, for more models with larger number of layers, a single epoch is insufficient. For instance, in ResNet-50 we achieve an  of 70.35% after one epoch, but we achieve 82.25% after 10 epochs. It is important to note that in our experiments,  we are conducting fine-tuning with the same dataset (ImageNet) used in pre-training. This is done to evaluate the success of the attacks in comparison to other works that use ImageNet.

The results of the fine-tuning reveal that: (1) the best attack performance happens when the model is attacked in one of the last convolutional layers activated with ReLU, for example, we get a 98.30% in  for one epoch in VGG-19 (2) our approach is potentially hard to caught in real deployments since the models are still highly accurate. For instance, in the VGG-19 experiments, the infected model achieved 73.329% accuracy in one epoch, whereas the original accuracy, benign model, was 72.376%, as shown in Table <ref>. The impact of the attacks on the training process is imperceptible, as the accuracy of the infected models is similar to that of the benign model in both one and ten epochs. Consequently, the attack itself would likely go unnoticed.







 §.§ Paintings Dataset


In <Ref>, we present the results using the Paintings dataset, which is not part of the original domain of the models. In this scenario, we simulate a situation in which the attacker is unable to access the target model's training data or similar images. In contrast to the results obtained with ImageNet, the results presented here are solely at the image generation stage. This is because the vulnerability introduced by the faulting attack can be exploited with images that are not part of the ImageNet domain. The results indicate that attacks conducted in deeper layers have the highest ,  a finding that is consistent with the results obtained using ImageNet. For instance, we achieved an ASR of 99.70% in VGG-19, and 93.87% in  DenseNet-121.








§ DISCUSSION



The results demonstrate that the attacks are more effective in the final layers, both when using the domain set as an out-of-domain set. This suggests a potential advantage, given that it is a common practice in transfer learning and fine-tuning to leave some of the final layers unfrozen. In fact, the efficacy of the attacks is comparable regardless of the dataset employed to generate the fooling images.  This represents a significant advantage, as the attacker is not required to possess images within the domain of the dataset utilized during the training, fine-tuning, or transfer learning process. Furthermore, the attack on one epoch demonstrates that  is both efficient and powerful, as the ReLU-skip attack does not have to be applied in numerous instances.

Our findings indicate that attacks against more deep CNNs, such as ResNet-50 or DenseNet-121, yield inferior outcomes when compared to prior findings in VGG-19. Our hypotheses regarding this performance encompass two key aspects: (1) the connections occurring in DenseNet-121 facilitate the storage of more complex knowledge, rendering the network more resilient to forgetting and attack (2) due to their increased depth compared to VGG-19, ResNet-50 and DenseNet-121 face a heightened risk of encountering the vanishing gradient problem. As indicated in <cit.>, a deep ReLU network will eventually die in probability as the network's depth approaches infinity. Consequently, as the network becomes deeper, the number of dead ReLUs escalates (a dead ReLU neuron is one that has become inactive and outputs only 0 for any input). ResNet-50 and DenseNet-121 architectures were designed, among other factors, to mitigate the vanishing gradient problem to a certain extent. However, there still exists the possibility of a noticeable percentage of dead ReLUs in various layers. Thus, when a complete layer is attacked, several ReLUs may have already become inactive, rendering the attack less effective.

Concerning the fooling images generation phase <Ref>, which entails the exploitation of the vulnerability, certain parameters can be adjusted to exploit the vulnerability optimally. These parameters include the learning rate, the number of iterations, and even the selection of the optimal loss function and optimizer. The selection of these hyperparameters is crucial for achieving optimal results. Following experimentation, a general configuration that proved effective for different target classes is presented in sections <Ref> and <Ref>. <Ref> illustrates the loss functions that were trialed in the generation of the fooling images and the impact of these loss functions on the , particularly in the ResNet-50 infected models.




Finally, we show in <Ref> a comprehensive benchmarking of  against related attack types. Particularly, for the same scenario, we show our performance in comparison to weights-oriented backdoor <cit.> and targeted evasion <cit.>. Overall,  achieves either a competitive or higher attack success rate while featuring advantages that render it more attractive and effective than other related attacks. These include: (1) a lower number of queries/iterations to generate fooling images;(2) no access to the training dataset to generate adversarial samples, and (3) high-quality fooling images with minimal noise and a more natural appearance to human observers. Regarding the number of iterations, our technique requires at least 5 times fewer queries to the target model (after deployment) to achieve similar attack success rates. On the other hand, one strong advantage of  is that it only requires one image (i.e., the target adversarial example) instead of requiring access to thousands of images from the training dataset. Thirdly,  generates adversarial examples where the trigger is imperceptible for humans. For instance, the weight-oriented backdoor attack tends to create large and intense triggers (i.e., colored pattern) taking 9.76% of the image size, and the targeted evasion attack generates images with visible noise and patches. In contrast,  generates adversarial examples where the noise is low and in most cases imperceptible. 




§ COUNTERMEASURES




Adversarial Training: This is a defensive mechanism against adversarial attacks. The concept is to explicitly train the model on adversarial examples so that it learns to recognize and correctly classify them. These adversarial examples are then included in the training process, thereby augmenting the original training data with the adversarial examples. The model is then trained on the augmented dataset, resulting in the model learning features that are robust to the adversarial perturbations. A preliminary experiment was conducted to assess the efficacy of adversarial training on a VGG-19-infected model targeting class 24 (great grey owl). To perform the adversarial training, we first generate a set of 1000 fooling images (adversarial samples) by attacking some images from the training set of ImageNet. Afterwards, we then fine-tune the model (i.e., VGG-19) including a set of adversarial examples (1000), in addition to a sample of 34,745 images from the ImageNet training set. Finally, we test the ASR on the test dataset to check if the adversarial training was effective. For VGG-19, we found that after re-training with this augmented dataset and testing with the set of roughly 3K fooling images, the  drops to approximately 5.33%.




§ RELATED WORK








Adversarial backdoor attacks.
Backdoor attacks[In the literature, one can also find the term trojan attack. These two terms can be used interchangeably.] on neural networks have been gaining attention in the adversarial learning community <cit.>.
While the first works focused on adding additional hardware components to plant the backdoor <cit.>, most of the follow-up works utilized changes in the training data to inject the backdoor <cit.> stealthily.
These attacks work by poisoning training samples to encode backdoor functionality during the training process <cit.>.
The backdoor is then activated by adding a trigger to the model input, which generally causes misclassification to a target class.
It was shown that this trigger can be hard to distinguish by the human eye <cit.>.

Fault-based backdoor attacks.
One of the challenges of the poisoning-based backdoor attack is that the adversary needs to be able to tamper with the training data, which might be impractical in many scenarios.
The answer to that are the fault-based backdoor attacks that modify the model parameters directly without changing the training data.

    
  * Inference phase attacks. These attacks change the model weight after the training.
    The first attack in this direction presented a generic methodology for backdooring CNNs via targeted weight perturbations <cit.>. 
    As perturbing weights has been previously used for misclassifying outputs <cit.>, this hardware attack method was utilized by Rakin et al. <cit.> to plant a backdoor at the bit-level.
    After that, several works adopted this method in various contexts <cit.>.
    Generally, weight-flipping attacks assume the attacker can use Rowhammer <cit.> – a bug in random access memories that allows bit-flipping of target bits by rapidly rewriting the surrounding area of the memory.
    
  * Training phase attacks. There are currently two directions in this sub-category, each of them having a single representative as of now.
    The first direction is similar to the aforementioned attacks in the sense that the attacker utilizes Rowhammer to flip the bits – but this time during the model training.
    DeepVenom is an instantiation of such an attack <cit.>.
    However, compared to DeepBaR, DeepVenom requires a visible large trigger image embedded in the original image and therefore is not imperceptible to human eyes.
    Our approach falls within the second direction – tampering with the processor instructions during the training process.
    The closest work to ours is an attack called FooBaR <cit.> which inserts a backdoor during the training phase of the model by faulting activation functions.
    However, the generation of triggers for FooBaR requires constraint solving, which is generally an NP-complete problem.
    The solving complexity grows exponentially with the number of parameters, and therefore, does not scale to large networks that are used for image classification tasks such as CIFAR10 or ImageNet.
    Our approach, on the other hand, generates the fooling images with gradient descent and thus can be used on arbitrarily large networks.
    Another disadvantage of FooBar is that, unlike in DeepBaR, the fooling images are visibly different from the original ones.


To put our work into the context of backdoor attacks, we can use a classification chart shown in Figure <ref>, inspired by <cit.> but with two additional categories: attack method, and attack phase.
The attack method differentiates between the poisoning attacks, detailed at the beginning of this section, and fault-based attacks, where  belongs.
Because of that category, we had to add the attack phase as well – while the poisoning attacks are always performed during the training, the bit-flip-based fault-based backdoors can also be injected during the inference.
In Figure <ref>, the categories where our attack belongs are highlighted in red color.






Defenses against backdoor attacks.
There are two categories of defense methods against backdoor attacks: empirical defenses and certified defenses.
Empirical defenses <cit.> are based on characteristics of existing attacks and provide good effectiveness against those.
However, as there is no theoretical guarantee with this type, it is possible to adapt attacks in a way to bypass them <cit.>.
Certified defenses, on the other hand, provide provable security guarantees under some assumptions <cit.>.
The performance of these is generally weaker compared to empirical defenses as it is hard to fulfill the given assumptions.

It is important to note that the existing defenses against backdoors protect against poisoning-based attacks, and therefore do not effectively protect against fault-based attacks <cit.>.

Defenses against fault-based attacks.
As the majority of fault-based attacks are weight-oriented, most of the defenses focus on protecting against bit flips in memory.
Techniques in this direction include checksums <cit.>, hashing <cit.>, encoding <cit.>, usage of binarized neural networks <cit.>, or in-DRAM swapping <cit.>.
A method called NeuroPots <cit.> utilizes a proactive honeypot approach to “lure” the attacker to flip certain bits in the network, making the detection and recovery more efficient.

To develop an appropriate countermeasure against ReLU-skip attacks, one would have to turn to techniques aimed at cryptographic implementations, as control-flow integrity protection is well-studied in this domain <cit.>.
One can use software methods such as instruction duplication/triplication to provide some level of protection <cit.> or replace a vulnerable instruction with a sequence of fault-tolerant ones <cit.>.
In hardware, redundant circuits can be used to some extent – however, it was shown that with precise equipment, one can simultaneously inject identical faults into these circuits <cit.>. 

Another class of countermeasures utilizes a separate device, usually a sensor <cit.>, that can detect local disturbances of a potentially harmful character.
For a more complete overview of techniques used for protecting cryptographic implementations, an interested reader is suggested to look at <cit.>.




§ CONCLUSIONS



In this work, we introduce a novel fault-based backdoor attack, , that allows attackers to trigger targeted misclassification for several image-based neural network architectures (such as ResNet-50, VGG-19, and DenseNet-121). Our technique relies on simple fault injection techniques that are applied during training (i.e., fine-tuning) stages. Overall, our approach achieves greater attack success rates than state-of-the-art targeted adversarial techniques without compromising the performance of the benign pre-trained model. Remarkably, our technique does not require full training of complex surrogate models to create further adversarial examples. Moreover, our approach requires 50 times fewer queries to the model than other backdoor attacks when creating adversarial examples. Additionally, we have presented a countermeasure to protect deployed systems against our proposed fault-based attacks.


IEEEtran




 Paper Title: Accelerating Full Waveform Inversion By Transfer Learning.





[
    *
    August 5, 2024
==================

1113



 



§ ABSTRACT



Full waveform inversion (FWI) is a powerful tool for reconstructing material fields based on sparsely measured data obtained by wave propagation. For specific problems, discretizing the material field with a neural network (NN) improves the robustness and reconstruction quality of the corresponding optimization problem. We call this method NN-based FWI. Starting from an initial guess, the weights of the NN are iteratively updated to fit the simulated wave signals to the sparsely measured data set. For gradient-based optimization, a suitable choice of the initial guess, i.e., a suitable NN weight initialization, is crucial for fast and robust convergence. 

In this paper, we introduce a novel transfer learning approach to further improve NN-based FWI. This approach leverages supervised pretraining to provide a better NN weight initialization, leading to faster convergence of the subsequent optimization problem. Moreover, the inversions yield physically more meaningful local minima. The network is pretrained to predict the unknown material field using the gradient information from the first iteration of conventional FWI. In our computational experiments on two-dimensional domains, the training data set consists of reference simulations with arbitrarily positioned elliptical voids of different shapes and orientations. We compare the performance of the proposed transfer learning NN-based FWI with three other methods: conventional FWI, NN-based FWI without pretraining and conventional FWI with an initial guess predicted from the pretrained NN. Our results show that transfer learning NN-based FWI outperforms the other methods in terms of convergence speed and reconstruction quality. 




Keywords: 
transfer learning, neural network, deep learning, adjoint optimization, Full waveform inversion






§ INTRODUCTION



Civil infrastructure is often deteriorated due to exposure to loads, temperature fluctuations, natural disasters, and potential faults during construction. Assessing the damage present in the infrastructure requires non-destructive testing (NDT) <cit.>. In recent years, attempts have been made to exploit the potential of full waveform inversion (FWI) in NDT. Originating in seismology <cit.>, FWI infers material models of the earth's interior by fitting simulated seismic waves to measured data sets from active or passive sources <cit.>. The approach is adopted in engineering applications <cit.> utilizing ultrasonic waves to reveal information about unknown defects such as voids, cracks, and inclusions. Extracting detailed information about damage location and severity using FWI is difficult, as the problem is ill-posed and computationally expensive <cit.>.


Deep learning-based approaches are investigated to enrich FWI, potentially offering a chance to overcome the challenges mentioned above. There are several methods to incorporate neural networks (NNs) into FWI, the most common being using NNs in a supervised manner. In this context, the NN learns to map time series signals, such as displacements or pressures at sensor locations within the domain, to the corresponding field of material coefficients, often referred to as distribution, which describes wave velocities or density <cit.>. Another approach uses NNs to enrich the signals with low-frequency data <cit.>. Kleman et al. <cit.> use an NN to predict the actual velocity distribution by inputting the velocity distribution after the first iteration of conventional FWI. A common challenge with supervised approaches is the need for large training data sets. To address this limitation, generative adversarial networks can be utilized <cit.>. However, due to the absence of an accuracy measure in the supervised learning approach, the resulting predictions are often unreliable, thereby limiting their practical applicability.


Another way to take advantage of deep learning methods in FWI is to use them with conventional FWI <cit.>. Herrmann et al. <cit.> show that inflating the optimization space by using a massively overparameterized NN to predict the coefficients of a piecewise constant material field can serve as an implicit regularization of the optimization problem. Solving the corresponding problem with the Adam optimizer results in fast convergence to physically reasonable local minima. One reason for this may be the inherent bias of NN to learn the low-frequency content of the data <cit.>. Therefore, the algorithm converges to an image that has fewer high-frequency artifacts. The gradients used for this case are calculated first by the continuous adjoint method, where the derivative of the cost function with respect to the coefficients of the piecewise constant material field is calculated. Then, the gradient of the cost function with respect to the NN parameters is determined by automatic differentiation (see the gradient computation in <Ref>). In the sequel, we will refer to this approach of using the NN as a parameterization of the material field as NN-based FWI.


It is possible to further improve NN-based FWI using transfer learning. Transfer learning is a technique in machine learning where an NN trained on one data set is used for a different but related learning task. By doing so, the NN starts the downstream task with weights already adjusted to recognize relevant features, significantly reducing the training time. Simply put, transfer learning uses the knowledge from previous data to accelerate learning for related tasks. Transfer learning has already been utilized in various fields, such as for the classification of text <cit.> and images <cit.>, machine translation <cit.> and robotics <cit.>. This approach can also be used in conjunction with NN-based Full Waveform Inversion (FWI), where pretraining the NN provides a suitable initial guess for the optimization task. The training data set for this includes wave simulations on samples with arbitrary defects. In a recent study by Muller et al. <cit.>, their NN was pretrained directly on sensor data to predict an appropriate initial velocity distribution. A similar approach is described in <cit.>.


In this paper, although a similar idea is used, a significant improvement in the pretraining of the NN is proposed. Instead of simply using raw sensor data to predict an initial material distribution as in <cit.>, the gradient computed by the adjoint method with respect to the material coefficients of the piecewise constant representation (obtained during the first iteration with a homogeneous material) is utilized. This simplifies the architecture of the NN because the gradient of the conventional FWI from the first iteration has the same shape as that of the domain. The gradient contains an approximation of the shape of the damage (<Ref>), which renders the learning task easier. Furthermore, using pretrained NN removes the dependency of the convergence of the NN on the initial weight initialization. The paper at hand thereby proposes a method for accelerating the NN-based FWI <cit.> by pretraining the NN. Using the pretrained NN for FWI is called transfer learning with NN-based FWI. The proposed method is compared with three other approaches: conventional FWI, NN-based FWI without pretraining, and conventional FWI with an initial guess from the pretrained NN.


The remainder of the paper is organized as follows: In <ref>, the background of conventional and NN-based FWI is discussed, including the proposed transfer learning approach. In <ref>, we demonstrate the superior performance of the transfer learning NN-based FWI with the help of several numerical examples. Finally, <ref> concludes the paper's main findings. 



§ METHODOLOGY
 




 §.§ Physical model


In the sequel, the scalar wave equation is used. It describes the propagation of the scalar wave field u(x,t) in space x and time t in a medium whose spatial distribution of the density is given by ρ(x). The spatial distribution of the wave speed is denoted as c(x), the external force as f(x,t). With this notation, the 2D scalar wave equation reads

    ρ(x) ü(x, t)-∇·(ρ(x) c(x)^2 ∇ u(x, t)) = f(x,t),   x=[x,y] ∈Ω, t ∈𝒯,

where Ω is the 2D spatial domain and 𝒯 = [0, T_max] the time interval of interest. Appendix A of <cit.> gives a detailed derivation of this equation. Bürchner et al. <cit.> show that using mono-parameter FWI to only invert for density leads to better identification of defects with strong contrasts such as voids. Therefore, a dimensionless scaling function γ(x) is introduced to parameterize the scalar wave equation and the following optimization problem. Assuming ρ_0 and c_0 correspond to undamaged background material, the unknown material field γ(x) scales the density, i.e., ρ(x) = γ(x)ρ_0, while the wave speed remains constant, i.e.,  c(x)=c_0. This leads to the density-scaled version of the scalar wave equation, which is used throughout the paper,


    γ(x)ρ_0ü(x,t)-∇·(γ(x) ρ_0 c_0^2 ∇ u(x, t)) = f(x,t),   x=[x,y] ∈Ω, t ∈𝒯.


The initial and boundary conditions are


    u(x, 0) =u̇(x, 0)=0     on Ω,



    ∂ u(x, t)/∂ x =0      on Γ_y×𝒯,


    ∂ u(x, t)/∂ y =0      on Γ_x×𝒯.





 §.§ Numerical configuration



In the paper, all investigations are performed with the same numerical configuration. The setup is shown in <Ref>.



Considering a 2D domain of 0.1×0.05, a finite differences (FD) solver is used to solve the corresponding wave equation (as discussed in detail in Appendix A of <cit.>). The domain is excited with a two-cycle sine burst at the positions x^s. Introducing the central frequency k_c with the corresponding angular frequency ω = 2 π k_c and amplitude A_0 (<Ref>), the external force reads

    f(x,t) =
        
        A_0sin(ω t)sin^2(ω t/2 n_c) δ(x - x^s)    for  0 ≤ t ≤2 π n_c/ω
    
            0    for 2 π n_c/ω < t

where n_c = 2 is the number of cycles and x^s is the source position of the underlying experiment.


The undamaged density and wave speed of the material is  2700kg/m^3 and  6000 m/s. The central excitation frequency is 500. The synthetic reference data is generated with simulations on a grid with 512 × 256 grid points and 2000 time steps with dt=3e-8. For the FWI, a different grid configuration has been used to avoid the “inverse crime” <cit.>. The domain is discretized with 256 × 128 grid points, and the time integration is performed with 1000 time steps and a time step length of dt=6e-8. The sources and sensors are positioned along the top edge of the domain. There are four sources in total and the simulation with the excitation from one sensor is equivalent to carrying one experiment. The sensors are distributed symmetrically about the center of the plate, 36 grid points apart. 24 sensors spaced 6 grid points apart are placed between the rightmost and leftmost sources. The excitation of one source is equivalent to conducting one experiment. 



 §.§ Conventional FWI


Using conventional FWI, the problem at hand is parameterized by a set of coefficients γ̂ representing the scaling function γ(x) at all finite difference grid points. Introducing a misfit functional ℒ(γ̂) the corresponding optimization problem is described by

    γ̂^* = γ̂argmin ℒ (γ̂),

where γ̂^* is the optimal solution. The misfit functional sums up the squared error at all N^r receiver positions x^r for all experiments N^s

    ℒ (γ̂; u^o) = 1/2∫_T∫_Ω∑_s=1^N_s∑_r=1^N_r[ (û_s(γ̂;x,t)-u_s^o(x^r,t))^2 δ(x-x^r)  ]dΩ dt,

where û_s(γ̂;x,t) denotes a single simulated wave field given a set of coefficients γ̂ and u_s^o(x^r,t) a observed wave field. Measurements across all experiments are summarized by u^o={u_s^o}_s=1^N_s. 




For the computation of the derivative of the cost function with respect to the material coefficients, the adjoint method is used. For a detailed derivation, see e.g. <cit.>. Given a wave field u, the derivative with respect to one coefficient γ̂_i is

    d ℒ/d γ̂_i = ∑_s=1^N_s∫_𝒯[-ρ_0 u̇_s,i^†u̇_s,i + ρ_0 c_0^2 ∇ u_s,i^†·∇ u_s,i] dt,

where the subindex i denotes the solution of the corresponding primal and adjoint solution at the grid point i. The derivatives at all grid points form the gradient is ∇_γℒ(γ̂; u^o) and can readily be used in gradient-based optimization.



 §.§ NN-based FWI




  §.§.§ Hybrid gradient computation


In the NN-based FWI, an NN parameterized by θ is used to predict the scaling function coefficients, i.e., γ̂(θ). To this end, consider an NN N_γ that predicts the material coefficients at all grid points: 

    γ̂(θ) = N_γ(θ;k),

where θ contains all parameters of the NN and k is the input as shown in <Ref>. Therefore, the optimization problem becomes

    θ^* = θargmin ℒ (γ̂(θ); u^o).

The gradient of the cost function with respect to the NN parameters ∇_θ can be computed with the chain rule

    ∇_θℒ(γ̂; u^o) = ∇_γ̂ℒ(γ̂; u^o)_ adjoint method·∇_θγ̂_automatic
    differentiation.

As described in detail in <cit.>, the first term ∇_γℒ(γ̂(θ); u^o) is obtained through the adjoint method, while the latter ∇_θγ̂(θ) with automatic differentiation. The NN used to predict the density scaling function γ (<Ref>) has a generator type architecture <cit.>. The details of the generator architecture can be found in the Appendix <Ref>. The input k to the generator network is a tensor of size 128×8×4 filled with random values sampled from a normal distribution with zero mean and variance one. The output is the spatial distribution of the density scaling function, a tensor of size 1 × 256 × 128 equivalent to the size of the FD grid. The Adam optimizer <cit.> available in the PyTorch library <cit.> is used to solve the optimization problem. The NN-based FWI method is summarized in <Ref>.






  §.§.§ Transfer learning NN-based FWI



Transfer learning NN-based FWI uses a pretrained U-Net (<Ref>) to parameterize the density scaling function γ. The knowledge from a training data set is transferred to provide good weight initialization of the NN for the downstream task of FWI.



Pretraining: In <cit.>, the pretraining was carried out using observed wavefields as inputs and the true density distribution as the output. This already improved NN-based FWI. Yet, further improvements are possible by using the adjoint gradient ∇_γℒ(γ̂_0; u^o) from the first iteration as the input where the material field γ̂_0 is equal to 1 throughout the entire domain. This simplifies the design of the NN architecture since the shape of the gradient has the identical dimension as the discretized domain of computation itself, thus allowing the choice of the same shape for input and output (<Ref>). The gradient also already contains an approximation of the damage. This causes the initial guess to be closer to the sought density scaling function γ^true. Therefore, the training data set is D = {∇_γℒ(γ̂_0; u^o_i),γ^true_i}_i=1^N where N is the number of training samples. The supervised pretraining is carried out using the mean squared error (MSE) over all N samples as loss function given as
 
    ℒ_D = 1/N∑_i=1^N‖γ^true_i-γ̂_i‖_2^2,

 where γ̂ is the output of the NN N_γ. For this purpose, the NN, as shown in <Ref>, is used. Note that a different architecture is used in transfer learning NN-based FWI as compared to NN-based FWI, since the input tensor for pretraining is the same size as the output tensor, which was not the case earlier. Supervised training aims to minimize the average loss over the entire training data set. Pretraining the U-Net in such a manner requires one full gradient computation of the first conventional FWI iteration. Therefore, preparing the dataset for pretraining is more expensive than the one carried out in <cit.> since one adjoint simulation has to be carried out in addition to every forward simulation. However, it leads to better results as demonstrated in <Ref> because it provides a better initialization of the NN. 


A U-Net-based convolutional NN architecture <cit.> is employed for pretraining. More details about this architecture are in the Appendix (<ref>). The U-Net employs skip connections (<Ref>) to address the vanishing gradient problem. As previously mentioned, the input tensor is the adjoint gradient from the first conventional FWI iteration ℒ(γ̂_0; u^o), and the output is the predicted density scaling function γ̂. The U-Net is trained on a data set containing one ellipsoid-shaped damage parameterized using the semi-major and minor axes a and b, the center x_c, y_c, and the rotation θ, see <Ref>. The mean squared loss function is used along with RMSprop <cit.> optimization to train the network. The training of the U-Net is carried out with 800 samples for 100 epochs using a batch size of 80 and requires around 240 seconds on a Nvidia Quadro RTX 8000. More details about the optimal number of epochs and samples used for pretraining can be found in Appendix <ref> and <ref>, respectively.




 
 
Once the U-Net is trained, it is used for the downstream task of FWI. The optimization then updates the weights of N_γ based on the FWI loss function <Ref>. The updated <Ref> is:

 
    γ̂(θ^p) = N_γ(θ^p;∇_γℒ(γ̂_0; u^o))
        .


In the following, the pretrained U-Net is used to perform NN-based FWI (<Ref>). The input to the pretrained network is the adjoint gradient from the first conventional FWI iteration. However, the loss function is changed to <Ref>, and the gradient used to update θ^p is used for the NN-based FWI outlined in <Ref>.



  §.§.§ Conventional FWI with initial guess



For comparison, the initial guess provided by the pretrained network is used with conventional FWI. Therefore, the prediction γ̂ from the pretrained network N_γ(θ^p;∇_γℒ(γ̂_0; u^o)) is used as starting model.



  §.§.§ Importance of NN initialization

Weight initialization for NNs is important to achieve fast convergence and good accuracy. Using the weight initialization suggested by Glorot in <cit.> is common practice. This avoids the vanishing gradient problem but does not form an optimal starting point for FWI, as the random weights lead to the noisy initialization patterns depicted in the upper left part of <Ref>. On the contrary, an expected outcome is that the material remains undamaged almost everywhere except for locally confined flaws. An approximation to undamaged material can be provided by additionally initializing the weights of the last layer using a standard normal distribution with a narrow standard deviation of 0.01 and setting a rather high bias of, e.g. three. This removes most of the noise from the initialization of the density scaling function and provides a better starting point for NN-based FWI. As a last step, a sigmoid function maps the values to the range between zero and one. Furthermore, setting the weight of the last layer causes the method to be less dependent on the initialization of the weights. <Ref> shows the effect of the weight initialization on the reconstruction of a damage case within 35 iterations. This weight initialization is used for the generator network for NN-based FWI (<Ref>) as well as for the U-Net employed for transfer learning NN-based FWI (<Ref>). In the case of transfer learning NN-based FWI, the weights of the last layer are initialized in this way during pretraining, which improves the convergence of the training process.





§ RESULTS AND DISCUSSION
 




 §.§ Damage cases


To compare all four methods, i.e., conventional FWI (<Ref>), NN-based FWI (<Ref>), conventional FWI with an initial guess from pretraining (<Ref>), and transfer learning NN-based FWI (<Ref>), four damage cases are used, see <Ref>. The aim is to quantify the methods with respect to their reconstruction quality and convergence speed. The red color represents the region where the density scaling function has a value of 1, and the blue color represents a void/damage, i.e., a value of 10^-5 the scaling function. 



The four case studies are presented to assess the performance of transfer learning NN-based FWI. Then, the four methods are compared quantitatively using an average of over 100 test cases. Apart from the four case studies, we also compare the methods quantitatively with an average performance over 100 test cases. These cases were not part of the training data set used to pre-train the U-Net. All methods were executed on an Nvidia Quadro RTX 8000, and each case required ∼100seconds to complete.





 §.§ Case 1


<Ref> compares the four methods for the first damage case. The test case is similar to but not a sample from the data used for the pretraining of the network described in <Ref>. Therefore, it is expected that the pretrained network will perform better than other methods for this case. The first column in <Ref> shows that conventional FWI recovers the damage within 10-15 iterations. However, the reconstruction suffers from artifacts, which only reduce slightly with increasing iterations. The second column shows the individual snapshots of NN-based FWI without pretraining, which recovers the damage within 10 iterations with good accuracy and without any artifacts in the reconstruction.
The third column shows the results of the corresponding iterations for conventional FWI with an initial guess from the pretrained network. As expected, providing a good initial guess improves the adjoint method's reconstruction and helps recover the damage faster than without pretraining. Finally, the fourth column shows the results from transfer learning NN-based FWI. The reconstruction of the damage is much faster; within 5 iterations, the reconstructed shape of the damage is very accurate, leading to the best results for this case.  



 §.§ Case 2


<Ref> shows the comparison of the four methods for a damage case containing a rectangular-shaped hole and a circular hole. Note that the training set does not include a square hole; only elliptical holes were included (<Ref>). The first column shows conventional FWI, which recovers the shape approximately within 20-25 iterations. However, there are strong artifacts in the area just above the rectangular damage. The NN-based FWI recovers the rectangular-shaped hole and the circular hole within 20 iterations with good accuracy. The conventional FWI with an initial guess from the pretrained network performs much better than without pretraining and achieves a good reconstruction within 20 iterations with less artifacts in the reconstructed image. Nevertheless, the circular hole is not reconstructed very well. Finally, the transfer learning NN-based FWI performs the best. It achieves an almost perfect reconstruction within 10 iterations without any artifacts. As compared to the NN-based FWI without pretraining, the reconstructed shape is close to the true shape, thus demonstrating that the pretrained network performs very well for cases slightly outside the distribution of data on which it is trained.





 §.§ Case 3


For the third case study (<Ref>), three holes of different sizes need to be reconstructed. This renders the task to be quite different from the training data. The conventional FWI can roughly recover the rightmost damage while failing to recover the damages on the left side within 35 iterations. There are a lot of artifacts in the reconstruction and it is difficult to infer the correct shape and location of the damages from the reconstruction. The NN-based FWI can recover one damage shape within 15 iterations but struggles to reconstruct the other two. However, in this case, the reconstruction is better than conventional FWI, although the reconstructed shape of the damages is not accurate. The conventional FWI with an initial guess from the pretrained NN, delivers a similar reconstruction of the damage without pretraining. By contrast, the transfer learning NN-based FWI can quickly recover all the locations of the holes with good accuracy within 20 iterations. It improves the reconstructed shape of the damages in the subsequent iterations. The damaged shape is recovered more accurately than other methods, and the artifacts are negligible. Therefore, the transfer learning NN-based FWI performs best among all the methods in this case.  




 §.§ Case 4: Is transfer learning NN-based FWI always the best?


In very few cases, the transfer learning NN-based FWI does not perform as well as the other methods. The initial guess provided by the network is too inaccurate in some cases where the sought damage is very different from those of the training data. One such example is shown in <Ref>. The figure depicts two intersecting rectangular damages with three circular holes with a decreasing radius located on a straight line. Furthermore, there are two additional circular holes on the left and the right sides, resulting in six different damages to be recovered. This is a difficult case for FWI in general, specifically for the pre-trained network, because the training data included holes with only a single damage. Additionally, the smaller circular holes are located in the bigger circular hole's shadow region, making it more challenging to reconstruct the second and third holes.
 

The conventional FWI can partially recover the rectangular-shaped damage and one hole on the side. The circular holes behind the rectangular damage are not recovered. The NN-based FWI delivers a fast reconstruction of the whole damaged region using only 25-30 iterations to provide high accuracy. In the case of the conventional FWI with an initial guess from the pretrained network, the initial guess is inaccurate. The recovered damages are partially accurate for the rectangular-shaped damages. The circular holes on each side are also partially recovered. For the transfer learning  NN-based FWI, the method could only recover the shape of the rectangular-shaped damage and missed all the other damages. Upon increasing the total number of iterations, it is observed that it can recover the circular holes on both sides after 55 iterations. Still, it cannot recover the two intersecting rectangular damages with three circular holes behind each other with a decreasing radius ever after 100 iterations. This is one of the few cases where the damage reconstruction is unsatisfactory. However, the NN-based FWI without pretraining performs the best in this case, clearly indicating the advantage of using NN to discretize the material field. 









 §.§ Comparison over 100 test cases



FWI is carried out for 100 test cases to compare the four methods quantitatively, each undergoing 35 iterations. The cost function and the mean squared error between the reconstructed and true density scaling function over the iterations are used as a measure of comparison. Results are depicted in <Ref>.





The cost function for the transfer learning NN-based FWI reduces the fastest till 20 iterations. After 35 iterations, the cost is comparable for NN-based FWI and the transfer learning NN-based FWI. However, the mean squared error is the lowest for transfer learning NN-based FWI, followed by NN-based FWI. The mean squared error at the start of the iteration is the same for conventional FWI with an initial guess from the pretrained NN and transfer learning NN-based FWI. However, the MSE for conventional FWI without and with initial guess from pretrained NN increases slightly over the iterations due to the artifacts in the reconstructed density scaling function. Conversely, the MSE consistently decreases for NN-based FWI and transfer learning NN-based FWI, indicating good reconstruction of the density scaling function with fewer artifacts.
 

In the examples shown here (as well as in <cit.>), the reconstruction of a void using NN-based FWI contains considerably fewer artifacts. This is because NNs contain an inherent spectral bias <cit.> and prefer to learn the low-frequency content of data. On the contrary, the frequencies of the artifacts in the conventional FWI reconstructions are rather high. Therefore, the reconstruction of the NN-based FWI contains a particular kind of high-frequency filter that surprisingly does not inhibit the formation of sharp changes from material to void and vice versa, and, yet, does not exhibit artifacts. 



 §.§ Can transfer learning FWI recover from bad initial guess?


When using a transfer learning framework, a potential issue arises if the pretrained network provides an incorrect initial guess. This can occur when the data given to the pretrained network is different from the training data. To this end, case 2 from <Ref> is revisited. To emulate a bad pretraining, an incorrect initial guess is set manually for Case 2 (<Ref>) with a subsequent inversion. In Figure <ref>, it is evident that the initial guess in the first iteration differs strongly from the target damage shape. Although the transfer learning FWI required roughly 70 iterations (compared to the 10 iterations with Transfer learning FWI and 25 iterations with NN-based FWI experienced in <Ref>), it accurately recovers the damage. Similar results were observed for many cases, including Case 1, whereas for some more complicated cases, such as Case 3 and Case 4, the recovered shape was not as good as before. In general, a greater number of iterations is needed to satisfactorily recover the damages when the initial guess is inaccurate. 






 §.§ Drawbacks and limitations

Transfer learning NN-based FWI also suffers from some drawbacks that are common to any machine learning-based method. It requires extra effort to tune all the hyperparameters involved in the process, such as the learning rate, the learning rate scheduler, gradient clipping, cost scaling, and the number of samples for pretraining. This can be a time-intensive process. 


As in any other machine learning-based approach, carefully selecting training data is crucial for pretraining the network. Data that is too similar leads to overfitting, with the consequence that the network will not generalize well for out-of-sample data. However, selecting training data can be challenging, particularly if the type of expected damage is unknown. As the domain gets bigger, the task of generating sensible training data for pretraining, as well as hyperparameter tuning, will become more challenging. This indicates potential for further improvement in the architecture of NN models, which are less sensitive to the choice of hyperparameters.




§ CONCLUSION


NN-based FWI <cit.> uses an NN to parameterize the density scaling function. It relies on wave signals emitted at sources and received by the sensors to reconstruct flaws. This paper demonstrates that it is possible to improve the NN-based FWI using transfer learning, as introduced in <cit.>. The transfer learning is achieved with a U-Net, which is pretrained in a supervised manner — mapping the gradient from the first iteration of conventional FWI to the true density scaling function. The pretraining alleviates the strong dependency of the NN on the initialization of the NN weights. The combination of the NN ansatz in FWI and transfer learning through supervised pretraining greatly improves the convergence and reconstruction quality over conventional FWI. This includes conventional FWI relying on the same initial guess obtained through the pretrained NN. 


In addition to evaluating the improvements compared to the original transfer-learning-based scheme presented <cit.>, the transfer learning NN-based FWI is thoroughly compared to conventional FWI both with and without using the same pretrained network to provide a good starting point. In all cases, the presented transfer learning NN-based FWI outperforms conventional FWI. It might be argued that appropriate regularization techniques <cit.> lead to similar results as the NN-based FWI for conventional FWI. These can, however, equally be applied to both the NN-based FWI and transfer learning FWI. Whether additional regularization boosts the performance of transfer learning FWI and if similar results are obtainable with conventional FWI is the subject of future study.



§ ACKNOWLEDGEMENTS

The authors gratefully acknowledge the funding provided by the Georg Nemetschek Institut (GNI) through the joint research project DeepMonitor, which finances Divya Shyam Singh and Qing Sun, as well as the Geothermal-Alliance Bavaria (GAB) by the Bavarian State Ministry of Science and the Arts (StMWK), which finances Leon Herrmann. Furthermore, we gratefully acknowledge funds received by the Deutsche Forschungsgemeinschaft under Grant no. KO 4570/1-1 and RA 624/29-1 which support Tim Bürchner. Felix Dietrich would like to acknowledge the funds received by the Deutsche Forschungsgemeinschaft - project no. 468830823.



§ DECLARATIONS

Conflict of interest No potential conflict of interest was reported by
the authors.



§ DATA AVAILABILITY

We provide a PyTorch implementation of all methods in <https://doi.org/10.5281/zenodo.13150916>.










§ NEURAL NETWORK ARCHITECTURES






 §.§ Generator network architecture


<Ref> shows the detailed architecture of the generator network used for the NN-based FWI method. The generator network upsamples the input tensor of size 128 × 8 × 4 through its layers to match the finite differences grid of the domain (1 × 256  × 128). 




 §.§ U-Net architecture





Since both the input and the output data are a 2D matrix, a U-Net-based convolutional NN architecture is employed; see <ref>. The U-Net <cit.> architecture consists of a blocked structure where each block consists of a Convolution layer, a Batch normalization layer, and an activation layer repeated twice, followed by a Max pooling layer. The input and output tensor is 256 × 128 in shape. The skip connections (<Ref>) are used to overcome the vanishing gradient problem.


A small parametric study was carried out to find the best parameters for pretraining the U-Net and the optimum number of samples and epochs required for pretraining. It's worth noting that the optimum performance for the complete method is the performance of the pretrained network for the downstream task, i.e., carrying out the NN-based FWI. After each pretraining task, the network was run for 10 test cases, and the average was compared to the performance of each pretrained network. It was found that the optimum performance was achieved by continuing the pretraining further after the validation error had stopped reducing for the pretraining task. A similar observation was reported in <cit.>.




 §.§ Hyperparameter tuning

The hyperparameters for each method are tuned separately using 35 iterations to maximize the individual performance (Table <ref>). A learning rate scheduler with a polynomial decay (β·epoch+1)^α was also used in this process with α=-0.5 and β=0.2.





  §.§.§ Optimal number of epochs


To find the optimum number of epochs for pretraining, the network was trained for 50, 100, 200, and 500 epochs. This was done with three different numbers of samples for pretraining, 50 samples, 200 samples, and 500 samples, to find out if the change in the number of samples is a factor in the optimum number of epochs for pretraining. The hyperparameters used for this parametric study are listed in <ref>.


 
 Hyperparameter     Value 

  Learning rate     8 × 10^-4 
 
 Gradient clipping     5 × 10^-5 

α     -0.5  

β     0.2 

Batch size     No. of samples / 10 

 
tableHyperparameters used for pretraining the U-Net







<Ref> shows the result of a parametric study for a U-Net pretrained using 800 samples and 50,100,200,300,500 epochs for pretraining averaged over 100 test cases. The graph also looks similar for 50 and 200 samples, with 100 epochs performing the best for carrying out transfer learning NN-based FWI. Increasing the number of epochs further reduces the accuracy. The initial guess of the network trained for 100 epochs was the best compared to other cases, and the error reduced much faster than in other cases. Therefore, 100 epochs have been used to pretrain the U-Net.



  §.§.§ Optimal number of samples
 
A parametric study was conducted to determine the optimal number of data sets used for pretraining the U-Net. This study aims to determine the effect of increasing the number of samples on the performance of the U-Net for NN-based FWI. The number of samples used for pretraining the U-Net is 200, 300, 400, 500 and 800. The hyperparameters used for pretraining the network were kept constant for each case. These hyperparameters are shown in <Ref>. <Ref> depicts the mean squared error at each iteration averaged over 100 test cases. Each curve represents the corresponding amount of samples used for pretraining. <Ref> demonstrates the clear pattern that more samples improve the performance of transfer learning NN-based FWI. As can be seen, the MSE error is the least for pretraining using 800 samples for the downstream task. Therefore, the network pretrained from 800 samples is used for final comparison with other methods. 














10
url<#>1urlprefixURL href#1#2#2 #1#1

park2000impedance
G. Park, H. H. Cudney, D. J. Inman,
  http://dx.doi.org/10.1061/(asce)1076-0342(2000)6:4(153)Impedance-based
  health monitoring of civil structural components, Journal of Infrastructure
  Systems 6 (4) (2000) 153–160.
https://doi.org/10.1061/(asce)1076-0342(2000)6:4(153)
  doi:10.1061/(asce)1076-0342(2000)6:4(153).
<http://dx.doi.org/10.1061/(asce)1076-0342(2000)6:4(153)>

lynch2007overview
J. P. Lynch, http://dx.doi.org/10.1098/rsta.2006.1932An overview of
  wireless structural health monitoring for civil structures, Philosophical
  Transactions of the Royal Society A: Mathematical, Physical and Engineering
  Sciences 365 (1851) (2006) 345–372.
https://doi.org/10.1098/rsta.2006.1932
  doi:10.1098/rsta.2006.1932.
<http://dx.doi.org/10.1098/rsta.2006.1932>

kim2007health
S. Kim, S. Pakzad, D. Culler, J. Demmel, G. Fenves, S. Glaser, M. Turon,
  http://dx.doi.org/10.1145/1236360.1236395Health monitoring of civil
  infrastructures using wireless sensor networks, in: Proceedings of the 6th
  international conference on Information processing in sensor networks - IPSN
  ’07, IPSN ’07, ACM Press, 2007.
https://doi.org/10.1145/1236360.1236395
  doi:10.1145/1236360.1236395.
<http://dx.doi.org/10.1145/1236360.1236395>

tarantola1984inversion
A. Tarantola, http://dx.doi.org/10.1190/1.1441754Inversion of seismic
  reflection data in the acoustic approximation, GEOPHYSICS 49 (8) (1984)
  1259–1266.
https://doi.org/10.1190/1.1441754
  doi:10.1190/1.1441754.
<http://dx.doi.org/10.1190/1.1441754>

fichtner2010full
A. Fichtner, http://dx.doi.org/10.1007/978-3-642-15807-0Full Seismic
  Waveform Modelling and Inversion, Springer Berlin Heidelberg, 2011.
https://doi.org/10.1007/978-3-642-15807-0
  doi:10.1007/978-3-642-15807-0.
<http://dx.doi.org/10.1007/978-3-642-15807-0>

menke2018geophysical
https://www.sciencedirect.com/science/article/pii/B9780128135556099949Copyright,
  in: W. Menke (Ed.), Geophysical Data Analysis (Fourth Edition), fourth
  edition Edition, Academic Press, 2018.
https://doi.org/https://doi.org/10.1016/B978-0-12-813555-6.09994-9
  doi:https://doi.org/10.1016/B978-0-12-813555-6.09994-9.
<https://www.sciencedirect.com/science/article/pii/B9780128135556099949>

seidl2017full
R. Seidl, E. Rank, http://dx.doi.org/10.1063/1.4974657Full waveform
  inversion for ultrasonic flaw identification, in: AIP Conference
  Proceedings, Author(s), 2017.
https://doi.org/10.1063/1.4974657
  doi:10.1063/1.4974657.
<http://dx.doi.org/10.1063/1.4974657>

He2021NumericalUF
J. He, J. Rao, J. D. Fleming, H. Nath Gharti, L. T. Nguyen, G. Morrison,
  http://dx.doi.org/10.1088/1361-665X/ac0f44Numerical ultrasonic full
  waveform inversion (fwi) for complex structures in coupled 2d solid/fluid
  media, Smart Materials and Structures 30 (8) (2021) 085044.
https://doi.org/10.1088/1361-665x/ac0f44
  doi:10.1088/1361-665x/ac0f44.
<http://dx.doi.org/10.1088/1361-665X/ac0f44>

burchner2023immersed
T. Bürchner, P. Kopp, S. Kollmannsberger, E. Rank,
  https://www.sciencedirect.com/science/article/pii/S0045782523000166Immersed
  boundary parametrizations for full waveform inversion, Computer Methods in
  Applied Mechanics and Engineering 406 (2023) 115893.
https://doi.org/https://doi.org/10.1016/j.cma.2023.115893
  doi:https://doi.org/10.1016/j.cma.2023.115893.
<https://www.sciencedirect.com/science/article/pii/S0045782523000166>

BKK23
T. Bürchner, P. Kopp, S. Kollmannsberger, E. Rank,
  https://www.sciencedirect.com/science/article/pii/S0045782523004103Isogeometric
  multi-resolution full waveform inversion based on the finite cell method,
  Computer Methods in Applied Mechanics and Engineering (2023) 116286https://doi.org/https://doi.org/10.1016/j.cma.2023.116286
  doi:https://doi.org/10.1016/j.cma.2023.116286.
<https://www.sciencedirect.com/science/article/pii/S0045782523004103>

virieux2009overview
J. Virieux, S. Operto, http://dx.doi.org/10.1190/1.3238367An overview
  of full-waveform inversion in exploration geophysics, GEOPHYSICS 74 (6)
  (2009) WCC1–WCC26.
https://doi.org/10.1190/1.3238367
  doi:10.1190/1.3238367.
<http://dx.doi.org/10.1190/1.3238367>

virieux2017introduction
J. Virieux, A. Asnaashari, R. Brossier, L. Métivier, A. Ribodetti, W. Zhou,
  http://dx.doi.org/10.1190/1.9781560803027.entry66. An introduction to
  full waveform inversion, Society of Exploration Geophysicists, 2014, pp.
  R1–1–R1–40.
https://doi.org/10.1190/1.9781560803027.entry6
  doi:10.1190/1.9781560803027.entry6.
<http://dx.doi.org/10.1190/1.9781560803027.entry6>

yang2019deep
F. Yang, J. Ma, http://dx.doi.org/10.1190/geo2018-0249.1Deep-learning
  inversion: A next-generation seismic velocity model building method,
  GEOPHYSICS 84 (4) (2019) R583–R599.
https://doi.org/10.1190/geo2018-0249.1
  doi:10.1190/geo2018-0249.1.
<http://dx.doi.org/10.1190/geo2018-0249.1>

wu2019inversionnet
Y. Wu, Y. Lin, Inversionnet: An efficient and accurate data-driven full
  waveform inversion, IEEE Transactions on Computational Imaging 6 (2020)
  419–433.
https://doi.org/10.1109/TCI.2019.2956866
  doi:10.1109/TCI.2019.2956866.

Wang2020VelocityMB
W. Wang, J. Ma, http://dx.doi.org/10.1190/geo2018-0591.1Velocity model
  building in a crosswell acquisition geometry with image-trained artificial
  neural networks, GEOPHYSICS 85 (2) (2020) U31–U46.
https://doi.org/10.1190/geo2018-0591.1
  doi:10.1190/geo2018-0591.1.
<http://dx.doi.org/10.1190/geo2018-0591.1>

rao2023quantitative
J. Rao, F. Yang, H. Mo, S. Kollmannsberger, E. Rank,
  https://www.sciencedirect.com/science/article/pii/S0022460X22006010Quantitative
  reconstruction of defects in multi-layered bonded composites using fully
  convolutional network-based ultrasonic inversion, Journal of Sound and
  Vibration 542 (2023) 117418.
https://doi.org/https://doi.org/10.1016/j.jsv.2022.117418
  doi:https://doi.org/10.1016/j.jsv.2022.117418.
<https://www.sciencedirect.com/science/article/pii/S0022460X22006010>

jin2018learn
Y. Jin, W. Hu, X. Wu, J. Chen,
  http://dx.doi.org/10.1190/segam2018-2997901.1Learn low-wavenumber
  information in fwi via deep inception-based convolutional networks, in: SEG
  Technical Program Expanded Abstracts 2018, Society of Exploration
  Geophysicists, 2018.
https://doi.org/10.1190/segam2018-2997901.1
  doi:10.1190/segam2018-2997901.1.
<http://dx.doi.org/10.1190/segam2018-2997901.1>

sun2020extrapolated
H. Sun, L. Demanet, https://doi.org/10.1190/geo2019-0195.1Extrapolated
  full-waveform inversion with deep learning, Geophysics 85 (3) (2020)
  R275–R288.
http://arxiv.org/abs/https://pubs.geoscienceworld.org/geophysics/article-pdf/85/3/R275/5018257/geo-2019-0195.1.pdf
  arXiv:https://pubs.geoscienceworld.org/geophysics/article-pdf/85/3/R275/5018257/geo-2019-0195.1.pdf,
  https://doi.org/10.1190/geo2019-0195.1
  doi:10.1190/geo2019-0195.1.
<https://doi.org/10.1190/geo2019-0195.1>

yang2022deep
T. Yang, X. Chai, W. Nie, J. Yu,
  https://library.seg.org/doi/abs/10.1190/iwmg2021-35.1Deep-learning-based
  low-frequency reconstruction for full-waveform inversion, 2022, pp.
  134–137.
http://arxiv.org/abs/https://library.seg.org/doi/pdf/10.1190/iwmg2021-35.1
  arXiv:https://library.seg.org/doi/pdf/10.1190/iwmg2021-35.1, https://doi.org/10.1190/iwmg2021-35.1 doi:10.1190/iwmg2021-35.1.
<https://library.seg.org/doi/abs/10.1190/iwmg2021-35.1>

li2022high
Y. Li, A. Alali, T. Alkhalifah,
  https://www.earthdoc.org/content/papers/10.3997/2214-4609.202210778High-wavenumber
  extrapolation for band-limited fwi results using deep learning 2022 (1)
  (2022) 1–5.
https://doi.org/https://doi.org/10.3997/2214-4609.202210778
  doi:https://doi.org/10.3997/2214-4609.202210778.
<https://www.earthdoc.org/content/papers/10.3997/2214-4609.202210778>

kleman2023full
C. Kleman, S. Anwar, Z. Liu, J. Gong, X. Zhu, A. Yunker, R. Kettimuthu, J. He,
  http://dx.doi.org/10.1115/1.4062092Full waveform inversion-based
  ultrasound computed tomography acceleration using two-dimensional
  convolutional neural networks, Journal of Nondestructive Evaluation,
  Diagnostics and Prognostics of Engineering Systems 6 (4) (Mar. 2023).
https://doi.org/10.1115/1.4062092
  doi:10.1115/1.4062092.
<http://dx.doi.org/10.1115/1.4062092>

araya2019deep
M. Araya-Polo, S. Farris, M. Florez,
  http://dx.doi.org/10.1190/tle38110872a1.1Deep learning-driven
  velocity model building workflow, The Leading Edge 38 (11) (2019)
  872a1–872a9.
https://doi.org/10.1190/tle38110872a1.1
  doi:10.1190/tle38110872a1.1.
<http://dx.doi.org/10.1190/tle38110872a1.1>

he2021reparameterized
Q. He, Y. Wang, http://dx.doi.org/10.1190/geo2019-0382.1Reparameterized
  full-waveform inversion using deep neural networks, GEOPHYSICS 86 (1) (2021)
  V1–V13.
https://doi.org/10.1190/geo2019-0382.1
  doi:10.1190/geo2019-0382.1.
<http://dx.doi.org/10.1190/geo2019-0382.1>

zhu2022integrating
W. Zhu, K. Xu, E. Darve, B. Biondi, G. C. Beroza,
  http://dx.doi.org/10.1190/geo2020-0933.1Integrating deep neural
  networks with full-waveform inversion: Reparameterization, regularization,
  and uncertainty quantification, GEOPHYSICS 87 (1) (2021) R93–R109.
https://doi.org/10.1190/geo2020-0933.1
  doi:10.1190/geo2020-0933.1.
<http://dx.doi.org/10.1190/geo2020-0933.1>

herrmann2023use
L. Herrmann, T. Bürchner, F. Dietrich, S. Kollmannsberger,
  http://dx.doi.org/10.1016/j.cma.2023.116278On the use of neural
  networks for full waveform inversion, Computer Methods in Applied Mechanics
  and Engineering 415 (2023) 116278.
https://doi.org/10.1016/j.cma.2023.116278
  doi:10.1016/j.cma.2023.116278.
<http://dx.doi.org/10.1016/j.cma.2023.116278>

Muller2023
A. P. O. Muller, J. C. Costa, C. R. Bom, M. Klatt, E. L. Faria, M. P.
  de Albuquerque, M. P. de Albuquerque,
  http://dx.doi.org/10.1093/gji/ggad215Deep pre-trained fwi: where
  supervised learning meets the physics-informed neural networks, Geophysical
  Journal International 235 (1) (2023) 119–134.
https://doi.org/10.1093/gji/ggad215
  doi:10.1093/gji/ggad215.
<http://dx.doi.org/10.1093/gji/ggad215>

rahaman2019spectral
N. Rahaman, A. Baratin, D. Arpit, F. Draxler, M. Lin, F. Hamprecht, Y. Bengio,
  A. Courville, https://proceedings.mlr.press/v97/rahaman19a.htmlOn the
  spectral bias of neural networks, in: K. Chaudhuri, R. Salakhutdinov (Eds.),
  Proceedings of the 36th International Conference on Machine Learning, Vol. 97
  of Proceedings of Machine Learning Research, PMLR, 2019, pp. 5301–5310.
<https://proceedings.mlr.press/v97/rahaman19a.html>

nigam2000text
K. Nigam, A. K. Mccallum, S. Thrun, T. Mitchell, Machine Learning 39 (2/3)
  (2000) 103–134.
https://doi.org/10.1023/a:1007692713085
  doi:10.1023/a:1007692713085,
  http://dx.doi.org/10.1023/a:1007692713085[link].
<http://dx.doi.org/10.1023/a:1007692713085>

Qasim2022
R. Qasim, W. H. Bangyal, M. A. Alqarni, A. Ali Almazroi,
  http://dx.doi.org/10.1155/2022/3498123A fine-tuned bert-based
  transfer learning approach for text classification, Journal of Healthcare
  Engineering 2022 (2022) 1–17.
https://doi.org/10.1155/2022/3498123
  doi:10.1155/2022/3498123.
<http://dx.doi.org/10.1155/2022/3498123>

shaha2018transfer
M. Shaha, M. Pawar, Transfer learning for image classification, in: 2018 Second
  International Conference on Electronics, Communication and Aerospace
  Technology (ICECA), 2018, pp. 656–660.
https://doi.org/10.1109/ICECA.2018.8474802
  doi:10.1109/ICECA.2018.8474802.

Hussain2018
M. Hussain, J. J. Bird, D. R. Faria,
  http://dx.doi.org/10.1007/978-3-319-97982-3_16A Study on CNN Transfer
  Learning for Image Classification, Springer International Publishing, 2018,
  p. 191–202.
https://doi.org/10.1007/978-3-319-97982-3_16
  doi:10.1007/978-3-319-97982-3_16.
<http://dx.doi.org/10.1007/978-3-319-97982-3_16>

Barret2018
B. Zoph, D. Yuret, J. May, K. Knight,
  http://arxiv.org/abs/1604.02201Transfer learning for low-resource
  neural machine translation, CoRR abs/1604.02201 (2016).
http://arxiv.org/abs/1604.02201 arXiv:1604.02201.
<http://arxiv.org/abs/1604.02201>

Surafel2018
S. M. Lakew, A. Erofeeva, M. Negri, M. Federico, M. Turchi,
  http://arxiv.org/abs/1811.01137Transfer learning in multilingual
  neural machine translation with dynamic vocabulary, CoRR abs/1811.01137
  (2018).
http://arxiv.org/abs/1811.01137 arXiv:1811.01137.
<http://arxiv.org/abs/1811.01137>

Tang2021
H. Tang, L. Notash, http://dx.doi.org/10.1115/DETC2021-69006Neural
  network based transfer learning for robot path generation, in: Volume 8B:
  45th Mechanisms and Robotics Conference (MR), IDETC-CIE2021, American Society
  of Mechanical Engineers, 2021.
https://doi.org/10.1115/detc2021-69006
  doi:10.1115/detc2021-69006.
<http://dx.doi.org/10.1115/DETC2021-69006>

Pereida2018
K. Pereida, M. K. Helwa, A. P. Schoellig, Data-efficient multirobot, multitask
  transfer learning for trajectory tracking, IEEE Robotics and Automation
  Letters 3 (2) (2018) 1260–1267.
https://doi.org/10.1109/LRA.2018.2795653
  doi:10.1109/LRA.2018.2795653.

Makondo2018
N. Makondo, B. Rosman, O. Hasegawa, Accelerating model learning with
  inter-robot knowledge transfer, in: 2018 IEEE International Conference on
  Robotics and Automation (ICRA), 2018, pp. 2417–2424.
https://doi.org/10.1109/ICRA.2018.8461218
  doi:10.1109/ICRA.2018.8461218.

kollmannsberger2023transfer
S. Kollmannsberger, D. Singh, L. Herrmann, Transfer learning enhanced full
  waveform inversion*, in: 2023 IEEE/ASME International Conference on Advanced
  Intelligent Mechatronics (AIM), 2023, pp. 866–871.
https://doi.org/10.1109/AIM46323.2023.10196158
  doi:10.1109/AIM46323.2023.10196158.

wirgin2004inverse
A. Wirgin, The inverse crime (2004).
http://arxiv.org/abs/math-ph/0401050
  arXiv:math-ph/0401050.

badrinarayanan2017segnet
V. Badrinarayanan, A. Kendall, R. Cipolla, Segnet: A deep convolutional
  encoder-decoder architecture for image segmentation, IEEE Transactions on
  Pattern Analysis and Machine Intelligence 39 (12) (2017) 2481–2495.
https://doi.org/10.1109/TPAMI.2016.2644615
  doi:10.1109/TPAMI.2016.2644615.

Kingma2014AdamAM
D. P. Kingma, J. Ba, https://arxiv.org/abs/1412.6980Adam: A method for
  stochastic optimization (2014).
https://doi.org/10.48550/ARXIV.1412.6980
  doi:10.48550/ARXIV.1412.6980.
<https://arxiv.org/abs/1412.6980>

NEURIPS2019_9015
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen,
  Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito,
  M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai,
  S. Chintala, Pytorch: An imperative style, high-performance deep learning
  library, in: Advances in Neural Information Processing Systems 32, Curran
  Associates, Inc., 2019, pp. 8024–8035,
  
<http://papers.neurips.cc/paper/9015-pytorch-an>
<-imperative-style-high-performance-deep-learning-library.pdf>.

Ronneberger2015
O. Ronneberger, P. Fischer, T. Brox,
  http://dx.doi.org/10.1007/978-3-319-24574-4_28U-Net: Convolutional
  Networks for Biomedical Image Segmentation, Springer International
  Publishing, 2015, p. 234–241.
https://doi.org/10.1007/978-3-319-24574-4_28
  doi:10.1007/978-3-319-24574-4_28.
<http://dx.doi.org/10.1007/978-3-319-24574-4_28>

Graves2013GeneratingSW
A. Graves, https://arxiv.org/abs/1308.0850Generating sequences with
  recurrent neural networks (2013).
https://doi.org/10.48550/ARXIV.1308.0850
  doi:10.48550/ARXIV.1308.0850.
<https://arxiv.org/abs/1308.0850>

glorot2010understanding
X. Glorot, Y. Bengio,
  https://proceedings.mlr.press/v9/glorot10a.htmlUnderstanding the
  difficulty of training deep feedforward neural networks, in: Y. W. Teh,
  M. Titterington (Eds.), Proceedings of the Thirteenth International
  Conference on Artificial Intelligence and Statistics, Vol. 9 of Proceedings
  of Machine Learning Research, PMLR, Chia Laguna Resort, Sardinia, Italy,
  2010, pp. 249–256.
<https://proceedings.mlr.press/v9/glorot10a.html>

Arridge2019
S. Arridge, P. Maass, O. Öktem, C.-B. Schönlieb,
  http://dx.doi.org/10.1017/S0962492919000059Solving inverse problems
  using data-driven models, Acta Numerica 28 (2019) 1–174.
https://doi.org/10.1017/s0962492919000059
  doi:10.1017/s0962492919000059.
<http://dx.doi.org/10.1017/S0962492919000059>

neyshabur2020being
B. Neyshabur, H. Sedghi, C. Zhang,
  https://proceedings.neurips.cc/paper_files/paper/2020/file/0607f4c705595b911a4f3e7a127b44e0-Paper.pdfWhat
  is being transferred in transfer learning?, in: H. Larochelle, M. Ranzato,
  R. Hadsell, M. Balcan, H. Lin (Eds.), Advances in Neural Information
  Processing Systems, Vol. 33, Curran Associates, Inc., 2020, pp. 512–523.
<https://proceedings.neurips.cc/paper_files/paper/2020/file/0607f4c705595b911a4f3e7a127b44e0-Paper.pdf>






 Paper Title: FSSC: Federated Learning of Transformer Neural Networks for Semantic Image Communication.


	
	
	
	 
	
FSSC: Federated Learning of Transformer Neural Networks for Semantic Image Communication
    
		
			Yuna Yan1, 
			Xin Zhang1, 
			Lixin Li1, 
			Wensheng Lin1, 
            Rui Li2,
            Wenchi Cheng3,
			and Zhu Han4 
		1School of Electronics and Information, Northwestern Polytechnical University, Xi’an, China, 710129
        2Samsung AI Center, Cambridge, UK
        3State Key Laboratory of Integrated Services Networks, Xidian University, Xi'an, China, 710071
		4Department of Electrical and Computer Engineering, University of Houston, Houston, TX, 77004
	
    August 5, 2024
==================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
In this paper, we address the problem of image semantic communication in a multi-user deployment scenario and propose a federated learning (FL) strategy for a Swin Transformer-based semantic communication system (FSSC). Firstly, we demonstrate that the adoption of a Swin Transformer for joint source-channel coding (JSCC) effectively extracts semantic information in the communication system. Next, the FL framework is introduced to collaboratively learn a global model by aggregating local model parameters, rather than directly sharing clients' data. This approach enhances user privacy protection and reduces the workload on the server or mobile edge.
Simulation evaluations indicate that our method outperforms the typical JSCC algorithm and traditional separate-based communication algorithms. Particularly after integrating local semantics, the global aggregation model has further increased the Peak Signal-to-Noise Ratio (PSNR) by more than 2dB, thoroughly proving the effectiveness of our algorithm. 




Semantic communication, federated learning, swin Transformer, privacy protection.




§ INTRODUCTION

Wireless communication is embracing an explosive growth of data traffic semantic nature, generated by various contemporary multi-media applications such as AR/VR, video streaming, gaming and telehealth. Traditional communication systems take the source messages (e.g. images and texts) and first encode these messages with source coding algorithms on the application layer, in order to remove redundant information. Subsequently, when the compressed information reaches the physical layer, it is then further channel-encoded by injection of redundant bits, to be resilient to the noise in the communication channel. There are three major factors that make this current approach inefficient in the 5th Generation Mobile Communication Technology(5G) and future generations of mobile networks. Firstly, although modern channel coding algorithm on its own has theoretically reached the Shannon limit on the physical layer <cit.>, when jointly considering the source coding and channel coding, the Shannon's optimality theory only holds when the message blocks are asymptotically infinitely long, which is impractical in the real-world <cit.>. Secondly, the separation-based approach<cit.> compresses and then upscales messages, thus inherently introducing extra latency, which lead to inferior performances for major time-sensitive 5G/6G and upcoming applications, e.g. tactile internet and autonomous driving. Third, channel coding strategies treat all bits of information equally, missing the opportunity to utilize the abundant semantic information available in the messages, especially in image and video contents.

To address these issues in the traditional communication regime, semantic communication<cit.> a.k.a. joint source and channel coding, has been proposed and studied, which focuses on the transmission of meaning rather than bits, helping to escape the “Shannon trap”<cit.>. 
Compared to conventional communication methods, semantic communication is centered around the relevant semantic features essential for transmitting information, thereby eliminating unnecessary redundancy. This approach leads to a significant reduction in data transmission latency and bandwidth consumption. To enhance the versatility of semantic communication systems, we propose the employment of a Swin Transformer-based semantic communication (STSC) system that primarily focuses on the reconstruction of global image signals. Its unique architecture enables the image-oriented semantic communication system to effectively capture deeper semantic features from potential semantic information and represent complex semantic relationships within the input data. In addition, the Swin Transformer is capable of reducing computational complexity and speeding up both model training and inference compared to traditional Transformer models. As a result, it is well-suited for handling large-scale semantic communication, especially in multi-user deployments.



However, image-oriented semantic communication targets applications that often involves users' private data, such as their faces, home interior, license plates, etc., which cannot always be uploaded to the cloud server. Moreover, most of the prior art on image-based semantic communication are point-to-point systems, which do not fit in many applications where multiple users are involved. In this paper, we propose a multi-user federated learning (FL) framework for semantic communication, showcasing the potential of efficient learning of a global semantic communication model that combines the learned knowledge from multiple users, without direct access to users' private data. We demonstrate that the proposed FL semantic system, when powered with a hierachical transformer architecture, i.e. the Swin Transformer, achieves a better image transmission quality at a faster convergence rate compared to the single-user.

In summary, the main contributions of this paper are summarized as follows:


	
  * We propose a Swin Transformer based semantic communication architecture, which aims at minimizing the difference between the transmitted and received images, so as to effectively communicate image semantic information.
	
  * In contrast to the centralised approach of sharing the user data directly to a cloud server in order to train a transformer neural network for semantic communication, we consider a cross-silo FL setting that employs multiple users to train locally their own model using local data, and share the local updates with a global server, thereby preserving users' data-privacy.
	
  * Through the use of well-established dataset for computer vision, simulation results show that our proposed framework surpasses the representative JSCC methods, and further boosts peak signal-to-noise ratio (PSNR) by over 2dB on top of aggregating local semantics.


The rest of this paper is organized as follows. In Section <ref>, we introduce the framework of the image semantic communication system based on FL and formulate the corresponding research problem. We then present the federated learning strategy for a Swin Transformer-based semantic communication (FSSC) in Section <ref>. In Section <ref>, we present numerical simulation results and showcase the semantic communication performance of FSSC as well as the convergence of the system. Finally, the paper is concluded in Section <ref>.




§ SYSTEM SETUP


In this paper, we consider a multi-user FL semantic communication system, as shown in Fig. <ref>, which consists of a cloud server and multiple clients. An FSSC client can be an edge server or a terminal device. When the client is an edge server, it is able to perform computational tasks, hence the training and inference of STSC happens on the client locally. When the client is a light-weight terminal device equipped with only limited or no computational power, the client should be served by a corresponding edge server which is trusted by the client to share data with in order to perform training. The federated training process is detailed as follows:

	
  * The cloud server initialize the STSC model for semantic communication, and distribute it to all participating clients. Each client will then perform n epochs of local training using its own data. Once the local training is completed, all clients transmit the updated model parameters to the cloud server, and the server will then aggregate the updated model parameters following Federated Averaging <cit.>, which completes one round of the federated learning. This process is repeated until the model converges;
	
  * During deployment of a trained STSC model, a user will initiate a communication request and sends a semantic message to the client. The client will then process this message by performing an inference using the trained STSC model with the message sample as the joint source-channel encoding process, and then transmit the output of the STSC model through the uplink channel to the receiver;
	
  * The transmitted message will then go through a noisy channel which typically add noisy corruption to the message. Once the receiver receives the message, it will then use the decoding part of the STSC model to recover the message. Note that the receiver in this setting could be either a cloud server, or a client.




Obviously, semantic information extraction and reconstruction of images are considered to be the key to the success of this system. Therefore, in this paper, the average mean square error (MSE) between the original input image and the reconstructed image is used as the loss function, which is defined as,


    MSE=1/L∑_i=1^n(x_i-x̂_i)^2,

where x_i is the original image vector, x̂_i is the reconstructed image vector, and L  represents the length of the image vector. A smaller MSE indicates that the reconstructed image is closer to the original image and the quality of image reconstruction is better.




§ IMAGE SEMANTIC ENCODER AND DECODER NETWORK


In this section, we first present an edge architecture for Swin Transformer-based semantic communication (STSC), which aims to capture long-term hierarchical semantic information from images. Additionally, to enhance the accuracy of semantic information extraction and ensure user privacy, we introduce the federated learning strategy for a Swin Transformer-based semantic communication (FSSC), which enables the training of edge models using FL across multiple devices hosted on cloud servers. As a result, our proposed scheme can effectively learn semantic representations from images transmitted by different users, and ensure reliable communication on various channels.



 §.§ STSC Model

Fig. <ref> shows the semantic communication network structure based on Swin Transformer. In this paper, joint source channel coding is adopted for semantic communication, in which the Swin Transformer module  <cit.> is adopted for encoding and decoding. Meanwhile, a non-trainable fully connected layer is used to simulate the physical channel.

At the transmitter, a set of N images 𝐗={x_i}_i=1^N is given, where x_i ∈ℝ^3 × H × W denotes the i-th image, and H and W denote height and width of the image, respectively. Firstly, the images are divided into non-overlapping patches through patch partion module. The size of patch is 4 × 4  and the vector dimension is converted to (H / 4,  W / 4,48). 
Then, the feature block sequence is linearly embedded. And through the learnable embedding matrix E, the feature blocks can be projected into an  embedding representation of arbitrary dimension C. After passing through the Swin Transformer block, the vector dimension is transformed to  (H / 4,  W / 4,C), 
and the whole process is called Stage 1. After the embedded representation, the feature blocks are fed into the Swin Transformer module, as shown in Fig. <ref>. As can be seen from the figure, the first patch merging layer and Swin Transformer block are combined into Stage 2. The patch merging layer concatenates each group of adjacent patches of size 2 × 2, so that the number of patch tokens becomes 1 / 4 of the original, i.e., H / 8×  W / 8. Meanwhile, the dimension of patch token is expanded by four times, i.e., 4C. In order to reduce the output dimension and realize the downsampling of the feature map, the patch merging layer performs a fully connected operation (implemented by a 1 × 1  convolutional layer). After this process, the dimension of the concatenated feature patch is reduced from 4C to 2C. Then, the patch goes through the swin Transformer block for feature transformation. Finally, the output dimension becomes  (H / 8,  W / 8,2C). The output x̃ passes through a fully connected layer to facilitate transmission in the channel. Specifically, it can be expressed as,

    s=W_1 x̃+b_1,

where, W_1 is the weight parameter matrix of the full connection, and  b_1 is the bias of the full connection.

In order to realize the joint training of encoder and decoder, a layer of fully connected neural network is used to simulate the channel. The neural network is actually a pair of input-output mappings, where the mapping relationship is determined by the neuron weight W_n and the bias b_1. Here, the weight W_n represents the channel gain, and the bias b_n is a random variable added to simulate the noise in the channel. The variance of the variable corresponds to the power of the channel noise, and its value mainly depends on the SNR and transmission power. After transmission in the physical channel, the signal s receives interference from channel noise and becomes the signal y when it reaches the receiver. 

At the receiver, the decoder consists of the Swin Transformer blocks. The function is to decode the signal y to restore the signal x̂.




 §.§ FL Training

In this paper, we develop the FSSC algorithm, by applying STSC to a FL framework, which can extract semantic information accurately and carry out semantic communication while protecting user privacy.While Fig. <ref> depics the overal architecture of FSSC, we detail below the specific training process of FSSC algorithm.

We use Fedavg<cit.> as the aggregation algorithm for our federated training. Assume that N clients participate in FL training. D_k  denotes the local Dataset of client E_k, and |D_k| is the size of the corresponding dataset. Thus, the training loss of the client based on the model in the local sample is

    Loss^k(ω)=1/|D_k|∑_i=1^|D_k|Loss_i^k(ω),

where  Loss^k(·) is the training Loss of the local client E_k, and Loss_i^k(·) is the loss of the sample x_i corresponding to the client. In the FSSC algorithm, Loss(·) is represented in (<ref>).

Therefore, the global federated training loss can be obtained by weighted average of the loss of each client according to the dataset size. The specific formula is as follows.


    Loss^∘(ω)=∑_k=1^N |D_k|/|D_o|Loss^k(ω),

where Loss^∘(ω) is the global loss after cloud aggregation, and |D_o| is the global dataset size. The purpose of the FL training is to find the global optimal model that minimizes the sum of training losses over all data. The specific federated training process is as follows.

Initialization. The dataset is non-uniformly split and distributed on each client. The semantic communication network based on Swin Transformer is initialized on the cloud server, and the initialized global model parameters W^o are broadcast to each client to build a local STSC network.

Model training. In the communication round t, the client participating in the training will conduct local training on STSC by using the local data set. And the stochastic gradient descent method is used to update the local model parameters. For each batch i, the process can be expressed as,


    ω_i+1^k ←ω_i^k-η∂Loss(ω_i^k ; b)/∂ω_i^k,

where  η is the learning rate. Each client saves its own trained local parameters and uploads them to the cloud server.

Parameter aggregation. The cloud server aggregates model parameters from different clients, and updates W_t^o to W_t+1^o . The aggregation mechanism shown in the following equation. Through this process, the update of the global shared model can be completed.


    w_t+1←∑_k=1^K |D_k|/|D_o| w_t+1^k,

where K is the total number of clients participating in training. This aggregation mechanism provides a certain guarantee for the stability of the system from the perspective of model convergence. Specifically, a certain degree of client-side fluctuations can be accepted without affecting the final convergence of the model.

Model convergence. The parameters of the global model W_t+1^o  are downloaded to each client. Then, the next round of training begins. Repeat this step until a specified global training cycle is reached or the model converges.



§ EXPERIMENT AND NUMERICAL RESULTS





 §.§ Simulation Settings

In this paper, we use the CIFAR-10 dataset<cit.> to evaluate the performance of the proposed FSSC algorithm. This dataset contains 50,000 training images and 10,000 test images, each with a size of 32 × 32. In addition, the number of clients is set to 3, and each client is assigned a different amount of non-overlapping training data. All clients use the same set of validation and test images. The training and testing environment are Windows 11+ CUDA 12.1, and the deep learning framework is Pytorch2.2.1.



In the experiments, the process of downloading and uploading model parameters between the client and the cloud server is regarded as an ideal situation. The focus is on the effects of global models obtained through local training and federated aggregation. The following introduces the initialization parameter settings for the semantic communication network based on Swin Transformer. Firstly, the network hyper-parameter C is set to 32. The network training all uses MSE as the loss function, the number of communication round is set to 60, the batch size of the data is 64, the learning rate is 0.001, and the Adam optimizer<cit.> is used. In all the tests, the compression ratio (CR) of the images is set to 0.33. 

To demonstrate the superiority of our method, we choose the typical JSCC <cit.> and the traditional communication algorithm as the comparison algorithm. The traditional communication algorithm adopts JPEG <cit.> as the source coding method, the channel coding adopts LDPC <cit.>, and the modulation mode and the demodulation mode are Quadrature Amplitude Modulation(QAM). And The modulation order is set to 4. 



 §.§ Simulation Results

In order to evaluate the semantic representation of the FSSC model, the experiments in this section compared the performance results of the local model algorithm STSC, JSCC and the traditional communication algorithms on AWGN channels, and verified the changes of the PSNR of the three methods under different SNR conditions, as shown in Fig. <ref>. It can be seen from Fig. <ref> that the PSNR of STSC and JSCC gradually increase with the increase of SNR. The STSC algorithm outperformed JSCC in terms of PSNR, with an average improvement of over 5 decibels. This significant advantage demonstrates the superior efficiency of the STSC algorithm in semantic feature extraction and reconstructing image data, resulting in a more accurate representation of the original image. At the same time, it can be noted that the PSNR of JPEG+LDPC+QAM remains unchanged when the channel conditions are very bad or greatly improved. This is because when the SNR is low, the traditional communication algorithm basically cannot transmit any semantic information. However, when the SNR is high, the PSNR reaches the performance saturation of the traditional communication algorithm. Thus, the accuracy of the proposed algorithm in this paper is significantly better than that of the traditional communication algorithms, especially when the SNR is improved, the performance improvement is very obvious. This is because while LDPC codes and QAM modulation can improve the robustness of data transmission, the JPEG and BPG compression algorithm are lossy and can lead to irreversible information loss. Combining the JPEG, LDPC and QAM techniques may reduce the error resilience of the overall system, leading to degraded image quality or data integrity in the presence of transmission errors. In comparison, the STSC algorithm retains the semantic information to the greatest extent, which ensures the effective transmission of images.



In order to evaluate the convergence of the proposed STSC model under the FL framework, the convergence analysis experiment of the FSSC model is also carried out in this section. Taking SNR=12 as an example, Fig. <ref> shows how the target loss varies as the training rounds/time increases in the case of FL training and local training, where the target loss is related to the degree of image reconstruction. As can be seen from the figure, on AWGN channels, the target loss of the communication model decreases rapidly in a short period of time, then fluctuates for a period of time, and finally remains unchanged. In other words, the model eventually reaches convergence. Moreover, the convergence efficiency of the FL training model is higher than that of the local training of the client. At the same time, the final training loss of global training model is relatively lower. This is because the local model uses a local single data set for local training, which makes it impossible to fully learn the global semantic distribution characteristics of the data. However, FL training has the advantages of data integration and improving model performance, and thus solves this problem well. 


Fig. <ref> shows how the PSNR change with the channel SNR for the global model and the local model without federated aggregation after model convergence. From the figure, it can be observed that with the increase of SNR, the PSNR of the global model and the local model without joint aggregation also significantly improved. In contrast to local models devoid of collaborative aggregation, global models can further enhance the PSNR by 2-3dB. This is attributed to the global model's ability to learn richer noise models and finer image semantic features through effective information aggregation and synchronization. Consequently, under high signal-to-noise ratio conditions, its PSNR can continue to improve, demonstrating superior generalization and noise resistance capabilities. This is because the image reconstruction is better as the channel condition improves. Moreover,  FSSC avoids the sharing of sensitive data, thus protecting the privacy of users and motivating users to participate in semantic communication under the FL framework to a certain extent.



§ CONCLUSION


In this paper, we proposed a semantic communication architecture for image transmission over wireless channels in multi-user cases, called FSSC. In this architecture, the local STSC semantic communication system model accurately extracts semantic information on the respective dataset and carries out the semantic communication model training. Then, joint training is performed through FL parameter aggregation to minimize the target loss function of the reconstructed image. Simulation results show that the algorithm can converge effectively. This approach is able to combine the semantic information exists in the diverse data from clients, and distributes the computational workload from edge to users, without compromising the user's privacy. Under the premise of privacy protection, the proposed algorithm can greatly reduce the amount of data required by each client while ensuring the fidelity in recovered images. Therefore, the FSSC proposed in this paper is a promising candidate scheme for multi-user image semantic communication system. 

00

ref1
K. Lu, Q. Zhou, R. Li, Z. Zhao, X. Chen, J. Wu, and H. Zhang, “Rethinking modern communication from semantic coding to semantic communication,” IEEE Wireless Communications, vol. 30, no. 1, pp. 158-164, February 2023.

ref2
E. Bourtsoulatze, D. Burth Kurka and D. Gündüz, “Deep Joint Source-Channel Coding for Wireless Image Transmission," IEEE Transactions on Cognitive Communications and Networking, vol. 5, no. 3, pp. 567-579, September 2019.

ref3
N. Farsad, M. Rao, and A. Goldsmith, “Deep learning for joint source-channel coding of text,” in 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP), Calgary, AB, Canada, April 2018. 

ref4
H. Xie, Z. Qin, G. Y. Li, and B. H. Juang, “Deep learning enabled semantic communication systems” IEEE Transactions on Signal Processing, vol. 69, pp. 2663–2675, April 2021. 

ref5
K. Choi, K. Tatwawadi, A. Grover, T. Weissman, and S. Ermon, “Neural joint source-channel coding,” in Proceedings of the 36th International Conference on Machine Learning (ICML), Long Beach, CA, United states, May 2019. 

ref6
J. Xu, B. Ai, W. Chen, A. Yang, P. Sun, and M. Rodrigues, “Wireless image transmission using deep source channel coding with attention modules,” IEEE Transactions on Circuits and Systems for Video Technology, vol. 32, no. 4, pp. 2315-2328, April 2022. 

ref7
Claude E Shannon, “A mathematical theory of communication,” The Bell system technical journal, vol. 27, no. 3, pp. 379–423, July 1948.

ref8
D. Gündüz, Z. Qin, I. E. Aguerri, H. S. Dhillon, Z. Yang, A. Yener, K. K. Wong, and C.-B. Chae, “Beyond transmitting bits: Context, semantics, and task-oriented communications,” IEEE Journal on Selected Areas in Communications, vol. 41, no. 1, pp. 5–41, November 2022.

ref9
M. Kountouris and N. Pappas, “Semantics-empowered communication for networked intelligent systems,” IEEE Communications Magazine, vol. 59, no. 6, pp. 96–102, June 2021.

ref10
W. Yang et al., “Semantic Communications for Future Internet: Fundamentals, Applications, and Challenges,” IEEE Communications Surveys & Tutorials, vol. 25, no. 1, pp. 213-250, November 2023.

ref11
Z. Liu et al., “Swin Transformer: Hierarchical Vision Transformer using Shifted Windows,” in IEEE/CVF International Conference on Computer Vision (ICCV), Montreal, QC, Canada, October 2021.

ref12
K. Yang, S. Wang, J. Dai, K. Tan, K. Niu and P. Zhang, “WITT: A Wireless Image Transmission Transformer for Semantic Communications,” in IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), Rhodes Island, Greece, June 2023.

ref13
H.B. Mcmahan, E. Moore, D. Ramage, S. Hampson, and B. A. y. Arcas, “Communication-Efficient Learning of Deep Networks from Decentralized Data,” in International Conference on Artificial Intelligence and Statistics (AISTATS), Fort Lauderdale, Florida, USA, February 2017.

ref14
A. Krizhevsky, “Learning multiple layers of features from tiny images,” Master's thesis, University of Tront, April 2009.

ref15
D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” arXiv preprint arXiv:1412.6980, December 2014.


ref16
G. K. Wallace, “The JPEG still picture compression standard,” IEEE Transactions on Consumer Electronics, vol. 38, no. 1, pp.xviii-xxxiv, February 1992.

ref17
R. Gallager, “Low-density parity-check codes,” IRE Transactions on Information Theory, vol. 8, no. 1, pp. 21–28, January 1962.








 Paper Title: VecAug: Unveiling Camouflaged Frauds with Cohort Augmentation for Enhanced Detection.













0009-0003-1493-8040

  National University of Singapore
  & Shopee Singapore
  

fxiao004@comp.nus.edu.sg



Corresponding author

  National University of Singapore 
  Singapore

shaofeng@comp.nus.edu.sg



  Zhejiang University
  Hangzhou, China

cg@zju.edu.cn



  University of Michigan
  Ann Arbor, USA

jag@umich.edu



  National University of Singapore
  Singapore

ooibc@comp.nus.edu.sg



  Beijing Institute of Technology
  Beijing, China

meihui_zhang@bit.edu.cn
































§ ABSTRACT 

Fraud detection presents a challenging task characterized by ever-evolving fraud patterns and scarce labeled data.

Existing methods predominantly rely on graph-based or sequence-based approaches.



While graph-based approaches connect users through shared entities to capture structural information, they remain vulnerable to fraudsters who can disrupt or manipulate these connections.


In contrast, sequence-based approaches analyze users' behavioral patterns, offering robustness against tampering but overlooking the interactions between similar users.


Inspired by cohort analysis in retention and healthcare, this paper introduces , a novel cohort-augmented learning framework that addresses these challenges by enhancing the representation learning of target users with personalized cohort information.

To this end, we first propose a vector burn-in technique for automatic cohort identification, which retrieves a task-specific cohort for each target user.

Then, to fully exploit the cohort information, we introduce an attentive cohort aggregation technique for augmenting target user representations.



To improve the robustness of such cohort augmentation, we also propose a novel label-aware cohort neighbor separation mechanism to distance negative cohort neighbors and calibrate the aggregated cohort information.
By integrating this cohort information with target user representations,  enhances the modeling capacity and generalization capabilities of the model to be augmented.
Our framework is flexible and can be seamlessly integrated with existing fraud detection models.

We deploy our framework on e-commerce platforms and evaluate it on three fraud detection datasets, and results show that  improves the detection performance of base models by up to 2.48% in AUC and 22.5% in R@P_0.9, outperforming state-of-the-art methods significantly.










<ccs2012>
   <concept>
       <concept_id>10010405.10003550</concept_id>
       <concept_desc>Applied computing Electronic commerce</concept_desc>
       <concept_significance>500</concept_significance>
       </concept>
   
   
   
   
   
 </ccs2012>


[500]Applied computing Electronic commerce














: Unveiling Camouflaged Frauds with Cohort Augmentation for Enhanced Detection
    Meihui Zhang
    ^1 Department of Mathematics, Politecnico di Milano, p.za Leonardo da Vinci 32, 20133 Milano, Italy

^2 IFP Energies nouvelles, 1 et 4 avenue de Bois-Préau, 92852 Rueil-Malmaison, France

^3 Department of Engineering Mechanics, KTH Royal Institute of Technology, Brinellvägen 8, 114 28 Stockholm, Sweden

====================================================================================================================================================================================================================================================================================================================





§ INTRODUCTION




Fraudulent transactions pose a serious threat to e-commerce platforms, as they cause economic losses, damage reputations, and harm user experiences <cit.>.
According to a recent report, fraud costs 2.9% of global e-commerce revenue in 2023 [<https://www.identiq.com/blog/ecommerce-fraud-prevention-report>].
This is especially concerning for the fast-growing Southeast Asian market, which is expected to have 140 million new consumers by 2030 [<https://www.linkedin.com/pulse/why-southeast-asia-expected-worlds-next-e-commerce-powerhouse/>]. 

However, detecting fraud in real-world scenarios remains challenging due to the dynamic nature of fraud patterns and the scarcity of labeled data <cit.>.

To address the sparse transaction data issue and identify potential fraudsters among new users, numerous fraud detection systems have been developed, leveraging advanced machine learning algorithms and integrating diverse data sources for comprehensive analysis and detection. These systems typically fall into two categories: graph-based approaches <cit.> and sequence-based methodologies <cit.>. 
Graph-based systems connect users based on shared identity features, such as devices, IPs, and phones, and use graph neural networks (GNNs) to capture the structural information. 
However, certain fraudsters employ sophisticated techniques, such as device information manipulation or VPN usage, to evade explicit connections.
They also imitate the behaviors of legitimate users, such as clicking on the same products or pages or using similar IPs, to camouflage themselves <cit.>.
These fraudsters are hard to identify using graph-based methods and lower the effectiveness of GNN-based approaches. 
In contrast, sequence-based systems analyze users’ behavioral sequences, which are more resilient to manipulation than individual features <cit.>.
Nonetheless, sequence-based models often overlook cross-user interactions between the target user and similar peers, thus failing to fully exploit the information from similar users.
Therefore, we need tailored approaches to connect well-camouflaged fraudsters with similar users and use their data for more accurate user representations.



In this paper, we present , a novel learning-based cohort augmentation framework for fraud detection using cohort information.
Our framework is inspired by cohort analysis, which is widely used in retention and healthcare applications to study the relationships and behaviors of users based on their activity patterns.
Unlike traditional cohort analysis methods that group users into different cohorts based on a set of traits or a specific event, our framework constructs implicit user cohorts based on the relevance of their behavioral sequences.
In this manner, we can identify potential fraudsters who exhibit similar fraudulent behaviors, regardless of their identity features, such as devices, IPs, and phones.

For example, Figure <ref> illustrates the behavioral sequences of three users, where User B and User C both search for the same products, make quick payments, and leave positive reviews.
They can be regarded as personalized cohort neighbors of each other and their information can be shared to enhance the homophily for fraud prediction. 
In contrast, User A performs product searches and comparisons before buying, behaving more like a normal user and should not be included in their cohort neighbors.



As illustrated in Figure <ref>,  has four phases: the vector burn-in phase, the cohort identification phase, the cohort-augmented training phase, and the prediction phase. 
In the vector burn-in phase, we use a novel technique to generate task-specific cohort vectors for all training samples and store them in a vector database ℰ_pool along with their fraud labels. 
This technique separates the optimization of embedding vectors from their use in the later cohort augmentation, which improves the training efficiency and effectiveness. 
In the cohort identification phase, we select K augmentation neighbors and K negative neighbors for each target user from the vector pool, based on their similarity and label discrepancy. 
The augmentation neighbors are used to enrich the representation of the target user, while the negative neighbors are used to create a contrastive learning objective. 
In the cohort-augmented training phase,  performs attentive neighbor aggregation to combine the cohort information from the augmentation neighbors, and label-aware neighbor separation to push away the representations of the negative neighbors. 
By doing so,  learns a more robust and discriminative model for fraud detection. In the prediction phase,  can rapidly retrieve the augmentation neighbors from the vector pool and use them to improve the prediction of the target user.
In summary, we make the following contributions:
 
    
  * We introduce , a pioneering cohort augmentation framework that improves the prediction performance of existing fraud detection methods by incorporating rich cohort neighborhood information and task-specific embedding vectors.
    
  * We develop a novel vector burn-in technique that identifies personalized cohorts based on behavioral sequences and detects fraudsters who evade connections by identity features.
    
  * We propose a novel label-aware neighbor separation technique that reduces noise from neighbors of opposite classes and improves cohort-augmented learning robustness.
    
  * We test  on the state-of-the-art deep learning models in fraud detection and observe substantial performance improvement.  enhances base models by up to 2.48% in AUC and 22.5% in R@P_0.9, with negligible computational overhead.
    









We organize this paper as follows. Section <ref> introduces the fraud detection tasks and the cohort augmentation concept. Section <ref> explains 's technical aspects and shows how personalized cohorts help detect camouflaged fraudsters. 
Section <ref> presents experimental results of integrating  with four base models to demonstrate its effectiveness. Section <ref> discusses the related work on fraud detection, and section <ref> concludes this paper.



§ PROBLEM DEFINITION



In this section, we first formulate the fraud detection task on user behavioral data.
Next, we introduce the idea of learning-based cohort augmentation for fraud detection, followed by the definition of augmentation neighbors and negative neighbors. 






 §.§ Fraud Detection




Fraud Detection on Behavioral Data.
For a target user u_i ∈𝒰, a common fraud behavior detection system takes the user's behavioral sequence as input: 𝐱_𝐢 = {s_0, s_1, ⋯, s_m-1}, where s_i ∈𝒮 represents user's action and m is the length of the sequence.
The fraud detection task over behavioral data aims to detect whether users engage in suspicious activities given their behavioral sequences.
Each user u_i is characterized by a corresponding behavioral sequence 𝐱_𝐢 and a ground-truth label y_i.
Identifying fake reviews is a specific instance of fraud behavior detection, where actions correspond to purchases of different products.





 §.§ Cohort Augmentation for Fraud Detection

In the cohort augmentation framework , we retrieve a personalized cohort for each user, aiming to exploit the implicit relationships between target users and their hidden connections.
Specifically, the personalized cohort comprises K augmentation neighbors, which are utilized to enhance target user representations, and K negative neighbors, which contribute to the creation of an effective contrastive loss for distinguishing neighbors with different labels from the target user.

Learning-based Cohort Augmentation.
Given the behavioral sequence of u_i in the e-commerce platform as 𝐱_i = {s_1, s_2, ⋯, s_n} and the corresponding label 𝐲_i,
the objective of cohort augmentation is to pinpoint similar fraudulent users controlled by the same entity as the target fraud users, and to identify additional normal users who resemble the target normal user, leveraging their information to enhance the predictive performance of the model.
We can formulate learning-based cohort augmentation as:



    𝐲_i^a = f(g(𝐱_i; 𝐖_g); 𝐞_i; E_i^a; 𝐖_f),


where f(·) is the augmented learning function parameterized by 𝐖_f; g(·) is the representation learning function of the base model with parameters 𝐖_g;
𝐞_i and E_i^a are the vector of the target user and the set of task-specific cohort augmentation neighbors in the vector burn-in space respectively, derived by a separately trained vector burn-in model;
𝐲_i^a represents the cohort-augmented prediction logit.
The main motivation for introducing the personalized learning-based cohorts is to leverage the supplementary information of similar samples E_i^a for enhancing the representation power and robustness of the representation learning function g(·).



Augmentation Neighbors.
The augmentation neighbors, denoted as 𝒩_aug(u_i, K), are K training samples that have similar fraud behaviors to the target sample u_i, as measured by the distance between their burn-in vectors. 

The augmentation neighbors can have the same or different labels as u_i (fraudulent or non-fraudulent). The augmentation neighbors help to improve the representation learning of the target sample by providing cohort information. The set of burn-in vectors of the augmentation neighbors is denoted as E_i^a.


Negative Neighbors.
The negative neighbors, denoted as 𝒩_neg(u_i, K), are K training samples that have opposite labels from the target sample u_i but have similar fraud behaviors, based on the distance between their burn-in vectors. The negative neighbors are used to improve the robustness and stability of cohort-augmented learning by minimizing the mutual information between the cohort data and the target user. The set of burn-in vectors of the negative neighbors is denoted as E_i^n.



§ METHODOLOGIES



This section introduces , a novel framework that aims to boost the performance of existing fraud detection methods by enhancing their discrimination capability. 
We describe each module of  and how they are trained to achieve this goal.







 §.§ Overall Framework







The illustration of  for enhanced fraud detection tasks in Figure <ref> presents the four main modules operating in distinct phases:
the vector burn-in phase, the cohort identification phase, the augmented training phase, and the prediction phase.
First, the vector burn-in module will generate fraud-specific vectors 𝐞_i via an auxiliary encoder g_aux(·) and store the resultant embedding vectors of all training samples in the vector database ℰ_pool.
For each user u_i,  identifies its K augmentation neighbor vectors E_i^a and K negative neighbor vectors E_i^n from ℰ_pool.
These vectors capture the relationship between users based on their fraud behaviors.
Next, in the augmented training phase as shown in Figure <ref>,  aggregates the information of E_i^a to augment the latent representation 𝐡_i of 𝐱_i obtained via the main encoder g(·).

To address camouflaged frauds,  applies a label-aware neighbor separation mechanism to push away the K negative neighbors in the augmented representation space, which reduces the influence of fraudsters who imitate normal users and improves detection accuracy.
Finally, in the prediction phase, only the augmentation neighbor vectors E_i^a will be retrieved to augment the target user representation 𝐡_i for the final prediction.
For clarity, the outputs of the auxiliary encoder and the main encoder are called vector and representation, respectively.




 §.§ Vector Burn-in Phase














To generate embedding vectors that can identify cohorts of fraudsters, we propose a novel vector burn-in technique that captures the fraud behaviors within each sample.
Specifically, we train an auxiliary encoder g_aux(·) and an output layer ϕ_aux using back-propagation in a supervised manner to obtain embedding vectors, i.e., 𝐞_i = g_aux(𝐱_i).
These vectors encode fraud-related information that is useful for subsequent cohort-augmented learning.
The auxiliary encoder can leverage existing fraud detection models, such as those proposed in <cit.> or traditional sequence learning methods in <cit.>, to obtain the fraud-specific embedding of samples and use a multilayer perceptron (MLP) as the output layer to generate prediction logits:



    𝐲_i = ϕ_aux(𝐞_i),   ϕ_aux:ℝ^n_d↦ℝ


where n_d is the dimension of the embedding vectors. 
Binary cross entropy (BCE) is typically adopted as the objective function for the fraud detection task and is optimized via gradient descent.
In this way, task-specific embedding vectors can be generated, and the corresponding latent embedding space is called vector burn-in space, which will not be optimized in the augmented training phase.
The novel vector burn-in technique effectively decouples cohort identification from the base model training process, eliminating early training noise, and ensuring more accurate cohort neighbors.






 §.§ Cohort Identification Phase

The embedding vectors of training samples, denoted as 𝐞_i, along with their corresponding labels 𝐲_i, are stored in a vector database ℰ_pool.
To identify the cohorts specific to fraudulent activities, the target sample is fed into g_aux(·), generating a vector query to search for K augmentation neighbors in ℰ_pool.
During this phase, the Euclidean distance metric is employed for retrieving the cohort neighbors.
However, cohort augmentation for fraud detection faces a key challenge: the potential degradation of prediction performance.

This can occur when incorporating information from well-camouflaged fraudsters with that of normal users, or when excessively integrating normal user behavior information to enhance the representation learning of fraudsters. 

To mitigate this challenge, we adopt a strategy that involves retrieving both K augmentation neighbors and K negative neighbors for each target user during training.
This strategy enables the model to discern the behavior disparities between well-camouflaged fraudsters and normal users.
We construct a tailored contrastive loss to achieve this goal, which is illustrated in Figure <ref> and explained in Section <ref>. 
Here, E_i^a and E_i^n represent the augmentation neighbor vectors and negative neighbor vectors, respectively.
It is worth noting that during inference for augmented learning, only the augmentation neighbors E_i^a are retrieved, as the labels of testing samples are not available to identify the negative neighbors E_i^n.




Comparison to Dynamic Nearest Neighbor Identification.
Conventional learning-based neighbor augmentation methods in <cit.> dynamically identify nearest neighbors and iteratively utilize augmented samples for model training. In contrast, our approach separates neighbor identification from base model training, eliminating early training noise and improving efficiency. This separation ensures more accurate and reliable nearest neighbors, resulting in more effective cohort-augmented learning. It enhances computational efficiency by focusing solely on base model training during the initial stages, leading to faster convergence and reduced training time.





 §.§ Cohort-Augmented Training Phase

To harness the potential of neighboring vector information and effectively augment the target sample modeling, we perform an attentive neighbor aggregation for the K augmentation neighbors and propose a label-aware neighbor separation mechanism to distance the K negative neighbors.
We first transform the augmentation neighbor vectors and the negative neighbor vectors to the augmented representation space, which is the new representation space derived from the main encoder g(·):


    H_i^a= 𝐖^augE_i^a + 𝐛^aug



where 𝐖^aug∈ℝ^n_d × n_d is the weight vector, 𝐛^aug∈ℝ^n_d is the bias term, and H_i^a is the representations of the augmentation neighbors.

Then we attentively aggregate information from the augmentation neighbors to the target sample using the following equations:


    α_i,k = Softmax(𝐡_i,k^TW^att𝐡_i),   𝐡_i,k∈H_i^a
    𝐡_i^a = ∑_k=1^Kα_i,k·𝐡_i,k


where 𝐡_i,k is the representation of the k-th augmentation neighbor in the augmented representation space, and
𝐡_i is the new representation of the target sample 𝐱_i obtained via the main encoder g(·), i.e., 𝐡_i = g(𝐱_i).
W^att∈ℝ^n_d × n_d is the bilinear attention weight matrix, and α_i,k corresponds to the attention weight for 𝐡_i,k.
𝐡_i^a and 𝐡_i will then be fused to generate the cohort-augmented logits in the final prediction module. 



    α_i,k = exp(𝐡_i,k^TW^att𝐡_i)/∑_j=1^Kexp(𝐡_i,j^T𝐖^att𝐡_i)



The incorporation of neighboring vector information can substantially enhance the model's learning capacity, but it may also compromise its disambiguation capability.
As illustrated in Figure <ref>, the neighboring samples for u_i may encompass samples from opposite classes.
To mitigate the adverse effects of noisy information from the augmentation neighbors, we build a contrastive objective to separate the negative neighbors in the augmented representation space.
Specifically, we adopt the same transformation layer to obtain representations H_i^n for the negative neighbors.
Subsequently, we perform a pairwise summation of the target user representation and the negative neighbors' representations. 
These negative-neighbor-augmented user representations 
are then fed into the main output layer to generate K negative-neighbor-augmented logits for supervised contrastive learning. The construction of the supervised contrastive loss will be elaborated in Section <ref>. 




Comparison to Vanilla Neighbor Augmentation.
Our method, combining supervised contrastive learning with nearest neighbor identification, effectively separates negative neighbors from target samples in the augmented representation space. This reduces the influence of noisy information from augmentation neighbors, leading to improved model performance and enhanced robustness to label noise.





 §.§ Prediction Phase

The target sample representation 𝐡_i and the corresponding augmentation neighbor representation 𝐡_i^a will be fed into an output layer to generate cohort-augmented logit:



    𝐲_i^a = ϕ_main^a(𝐡_i + 𝐡_i^a),   ϕ_main^a:ℝ^n_d↦ℝ


We fuse the two representations via a summation, which empirically is more efficient and stable than concatenation.
The output layer ϕ_main^a will capture the task-specific information from the target sample and augmentation neighbors.
Meanwhile, we can obtain the negative-neighbor-augmented logits Y_i^n∈ℝ^n_p × K by feeding the sum of target user representation and the K negative neighbors' representations into the main output layer.




 §.§ Multi-Objective Joint Training


We use a joint training scheme for  that combines different objectives in one framework for effective fraud detection.
This helps  to capture intricate fraud behaviors and improve robustness against camouflaged frauds. 



Task #1: Main Supervised Prediction.
For a fraud detection task, we use BCE as a main objective function:

    ℒ^a_main = -1/N∑_𝒯_train𝐲_i log(σ (𝐲_i^a))+(1 - 𝐲_i)log(1-σ (𝐲_i^a))

where 𝐲_i and 𝐲_i^a are the ground truth label and cohort-augmented prediction logit respectively. N is the number of training samples, and σ(·) is the sigmoid function.
With this objective function,  can be trained effectively via gradient-based optimizers.



Task #2: Supervised Contrastive Learning.
We introduce a label-aware neighbor separation mechanism to counter the negative effect of noisy information from augmentation neighbors.
For this objective function, any negative-neighbor-augmented logit 𝐲_i,k^n in Y_i^n should incur greater losses than 𝐲_i^a, considering the ground truth 𝐲_i of the target sample. Drawing inspiration from prior works <cit.>, we construct the contrastive loss on the prediction logits rather than the cohort-augmented representations, as presented below:




    ℒ_sccl = -1/N∑_𝒯_trainlogexp(-(𝐲_i, 𝐲_i^a)/τ)/exp(-(𝐲_i, 𝐲_i^a)/τ) + ∑_k=1^Kexp(-(𝐲_i,𝐲_i,k^n)/τ)



where (𝐲_i, 𝐲_i^a) represents the distance of the ground-truth logits and cohort-augmented logits, and τ is the temperature hyperparameter.
With low temperatures, the loss is dominated by closer cohort-augmented logits.
This method enables supervised contrastive learning and prevents the cohort augmented representation from diverging too much from the original burn-in vector. In this manner, the personalized cohort from the burn-in phase remains useful in the augmentation phase.




Task #3: Latent Space Alignment.

Since cohort identification depends on vector distances between the target sample and candidates, it may deviate slightly from distances in the augmented representation space, potentially reducing the relevance of neighbor information to the current prediction task. To rectify this inconsistency, we introduce auxiliary vector regularization for consistent distances in both vector burn-in and cohort-augmented learning phases:



    ℒ_align= 1/N∑_𝒯_train𝐞_i - 𝐡_i^2


The final objective function of  is the summation of the above three tasks and the canonical weight decay regularization:



    ℒ= ℒ^a_main + αℒ_sccl + βℒ_align + λθ^2



where α and β are the hyperparameters that control the weights of the contrastive loss and auxiliary reconstruction loss respectively,
θ is the set of all model parameters,
and λ is the weight regularization strength.





Prediction Phase.
In the final prediction phase,  first obtains the burn-in vector of the target sample 𝐞_i as a query to retrieve the augmentation neighbors.
Next, the augmentation neighbor vectors E_i^a are transformed into the augmented representation space H_i^a and fused with the new representations 𝐡_i of the target sample.
Finally, the augmented representation is fed into the output layer to predict the fraud probability.


§ EXPERIMENTS


In our study, we integrate four established fraud detection models with  to assess the effect of the proposed label-aware cohort augmentation on detecting hidden fraudsters.
Particularly, our experiments are devised to answer the key research questions (RQs) as follows:
 

  * RQ1: Can  improve the existing sequence-based fraud detection model's performance without harming the legitimate users' activities?

  * RQ2: How do cohort augmentation and label-aware neighbor separation mechanisms contribute to fraud detection tasks?

  * RQ3: How does  perform as compared with existing neighbor augmentation approaches? 

  * RQ4: How does the size of augmentation and negative neighbor sets impact 's performance in fraud detection?

  * RQ5: How effective is  at unveiling fraud camouflage in graph-based fraud detection methods?




 §.§ Experimental Setup



  §.§.§ Datasets.

We mainly use two private and one public datasets to evaluate our proposed method. The private datasets are obtained from Shopee [<https://en.wikipedia.org/wiki/Shopee>], the largest online e-commerce platform in Southeast Asia. 

Each user is assigned a risk score by a fraud detection model based on their daily activities or purchasing history.
The public dataset is the Amazon [<https://en.wikipedia.org/wiki/Amazon_(company)>] dataset <cit.>, which is widely used for fake review detection. 
Fake review detection is a specific case of fraudulent behavior detection, where each action is a product review.
We compare the performance of the base models with and without  on the three datasets to demonstrate its effectiveness.



  §.§.§ Shopee Dataset
 The large-scale industrial datasets collected from Shopee adhere strictly to security and privacy protocols.
These datasets consist of historical activity logs originating from the same geographical regions, albeit varying in size, designated as Shopee-large and Shopee-small, spanning a specified timeframe.
Table <ref> provides an overview of the dataset statistics. 
For fraud behavior detection, each user is represented by 300 records, each containing details such as actions taken, timestamps, and associated attributes. These actions encompass a spectrum of user interactions within the Shopee platform, including "visit-homepage", "refund-request", and "password-reset-request". 
The fraud rate in the table is our sampled fraud rate for the experiments, not the actual fraud rate on Shopee.

We also collected a user-device interaction graph from Shopee, comprising 100k users and 54k devices. This graph dataset serves to validate the effectiveness of  in uncovering fraud camouflage. 
We employ two graph neural networks for fraud detection on this dataset.
Subsequently, we integrate these GNNs with  to observe performance improvements.
To simulate fraudsters' camouflage strategies, we sub-sample the Shopee graph by removing 3758 edges between certain fraudulent accounts and devices. 
Finally, we compare the performance of the two GNNs and their integration with  on this sub-sampled Shopee graph.




  §.§.§ Amazon Dataset

The Amazon dataset comprises diverse product reviews from categories such as electronics, books, CDs, and movies[<https://jmcauley.ucsd.edu/data/amazon/>]. Extensive use case analyses conducted in previous studies <cit.> have revealed patterns of fraudulent activities, including repeated purchases of identical items and the rapid submission of reviews. 
Utilizing the filtering algorithms in <cit.>, users are labeled based on their helpful vote percentages, with those below 20% as fake users and those above 80% as normal users.





  §.§.§ Baseline Methods

We integrate  with four sequence-based fraud detection models and compare its performance to two existing neighbor augmented learning methods.
Additionally, we utilize two graph neural networks to validate the effectiveness of  in uncovering camouflaged fraud.
Ablation tests are also conducted to show the efficacy of the techniques proposed in .



Sequence-based fraud detection base models. 
(1) BiLSTM <cit.> applies bidirectional LSTM to capture the dependency information among sequential data for fraud detection.
(2) LIC Tree-LSTM <cit.> builds a behavior tree for the user's behavioral sequence and studies fraud transaction detection with local intention calibration.
(3) HEN <cit.> uses field-level extractor and action-level extractor to hierarchically learn users' intentions and identify high-risk behavior sequences.
(4) MINT <cit.> employs a time-aware behavior graph to model user's behavioral data and utilizes a multi-view graph neural network to detect fraudulent behaviors.



Graph-based fraud detection base models. 
(1) GraphSAGE <cit.> performs inductive learning on graphs by sampling and aggregating node representations.
(2) GeniePath <cit.> integrates LSTMs into GCN layers to adaptively aggregate neighborhood information.



Existing neighbor augmentation methods. (1) GRASP <cit.> divides samples into different clusters and uses the cluster representations to enrich sample representations. (2) CSRM <cit.> dynamically retrieves different neighbors' information during model training for augmented learning.



Ablations. _BI is a variant of  that adopts a vanilla cohort augmentation architecture with a vector burn-in mechanism. _LA is a variant of  that only introduces the label-aware neighbor separation mechanism.




  §.§.§ Implementation Details

In our implementation, we utilize PyTorch for all models and Milvus <cit.> to store embedding vectors.



For the datasets, the split ratio of the training/validation/testing set is 60%/20%/20%.
We employ the Adam optimizer and initialize model parameters with Xavier initialization.
Early stopping is applied if the validation performance does not improve for 20 epochs.
The settings of the fraud detection task follow the practice in <cit.>. 
The embedding dimension of users' behavioral action is set to 64, the batch size to 128, the learning rate to 0.0001, and the maximum behavioral sequence length to 300.
The regularization strength is set to 1e-5.


The number of nearest neighbors K is also set to five for all models and datasets. 

We used grid search to optimize three key hyperparameters for  with different base models: the contrastive loss weight α, the auxiliary reconstruction loss weight β, and the temperature τ.
These hyperparameters control the importance of the contrastive objective, the balance between cohort-augmented learning and latent space alignment, and the focus on close or far-away cohort neighbors, respectively. 

We search α in (0.00001, 0.0001, 0.001, 0.01, 0.1, 1.0), β in (0.0000001, 0.000001, 0.00001, 0.0001, 0.001), and τ in (0.1, 0.3, 1.0, 3.0, 10.0).
These ranges were chosen to cover a wide spectrum of possible values for different fraud detection models. We found that τ = 1 was the best value for all models.
Table <ref> shows the values of α and β selected for  with various base models and datasets.







  §.§.§ Evaluation Metrics

For fraud detection tasks, we evaluate all methods with two widely used metrics: AUC (Area Under ROC) and R@P_k. AUC signifies the probability that positive cases' predictions are ranked ahead of negative cases. R@P_k represents the recall rate when precision equals k. 
In our fraud detection scenario, the goal is to identify as many suspicious entities as possible while not disrupting the regular operations of benign entities. We set k to 0.9 for all experiments. 
A superior cohort-augmented learning framework is expected to achieve higher AUC values than base models and notably enhance the recall rate over base models when precision values are set to 0.9.





 §.§ Effectiveness of Cohort Augmentation(RQ1)

Table <ref> shows the fraud detection performance of four state-of-the-art baselines with and without .
 improves both the AUC and R@P_0.9 of the base models significantly.
On the Shopee datasets, the AUC increases by at least 0.89% and R@P_0.9 by more than 4.31%. On the Amazon dataset, the AUC rises by at least 1.22% and R@P_0.9 by 5.81%. 

Notably, the substantial enhancement in R@P_0.9 indicates that  aids in identifying more fraudsters who escape the detection of base models.
These fraudsters mimic the behavior of normal users to avoid being caught.
 uses the task-specific embedding vectors of the base models to create personalized implicit user cohorts, linking camouflaged fraudsters implicitly to other users who are controlled by the same individuals or groups.
Unlike conventional GNN-based approaches vulnerable to intentionally masked identity features of fraudsters, our cohort identification methods rely on learned user behavioral representations to detect the camouflaged fraudsters effectively.

For example, in Shopee, users with similar behavioral sequences may have similar user representations learned by base models. 
However, the base models cannot adequately separate fraud from normal users due to the presence of many camouflaged fraudsters.

 connects camouflaged fraudsters to labeled fraud users and separates them from similar normal users using a label-aware neighbor separation mechanism, which enables better discrimination between suspicious and normal behaviors.





 §.§ Ablation Studies (RQ2)


We conduct ablation studies on  to demonstrate the impact of vector burn-in cohort augmentation and label-aware neighbor separation on its performance.



  §.§.§ Vector Burn-in for Cohort Augmentation

Figure <ref> demonstrates that the vector burn-in technique alone substantially improves base model performance.
This indicates that the learning-based cohort approach is effective for fraud detection. 
By constructing event-based networks for each user, we can leverage the similarities among users to infer their labels. 
Moreover, we can link unseen suspicious activities to existing samples, and obtain more relevant information for prediction. 
Fraud detection is a binary task, but fraud behavioral patterns are diverse and complex. 
The user behavioral representations learned from any base models will not form two clear clusters. 
Instead, they will form many clusters of varying sizes, especially for user accounts operated by the same person or group. 
Normal users’ representations are usually more distinct from each other. Therefore, if we construct a personalized cohort for each user, it will mainly affect fraud users and make their clusters more compact. 
Fraud users not affiliated with any cluster are connected to other fraud users who are similar to them. 
This has minimal impact on normal users, who are less likely to be operated by a common entity, and whose cohort-augmented representation will not change much. As a result, we can detect more fraudulent users without compromising the performance of normal users.




  §.§.§ Label-aware Neighbor Separation Mechanism
  
Figure <ref> demonstrates a significant enhancement in both AUC and R@P_0.9 metrics across both datasets upon integrating MINT with _LA. 
While cohort augmentation fosters connections between various fraudsters and enhances learned representations, it may inadvertently link adept fraudsters with normal users exhibiting similar behavioral patterns. 
This scenario introduces considerable noise through direct message passing between fraudulent and normal users, thereby notably compromising model performance. 
To mitigate this issue, the label-aware neighbor separation mechanism deliberately distances these challenging negative neighbors during cohort augmentation via a supervised contrastive loss. 
Given the original base models' inability to discern target users from their difficult negative counterparts, the contrastive objective function serves to recalibrate the learning bias, prioritizing the treatment of these negative pairs and effectively separating them within the augmented learning space. 
This calibration is achieved through the construction of the contrastive loss on the augmented prediction logits, primarily influencing the output layers while also exerting a lesser influence on the encoder training.



 §.§ Comparison with Existing Neighbor Augmentation Methods (RQ3)

In this subsection, we compare the effectiveness of  with two existing methods that use neighborhood information to enhance target users: GRASP <cit.> and CSRM <cit.>. 
Figure <ref> illustrates the performance comparisons. 
GRASP clusters the samples into groups based on the task-specific vectors using K-Means, and uses the centroid vectors as group representations. 
It constructs a KNN graph for different groups, with each user aggregating information from its group and other connected groups. 
However, GRASP functions as a three-layer graph neural network and passes excessive information from the neighbors to the target users, leading to severe over-smoothing. 
Moreover, GRASP compromises users’ personalized behavioral patterns and performs worse than the base model on both datasets. 
CSRM dynamically retrieves the nearest neighbors for each user, demonstrating more effective neighbor augmentation than GRASP. 
However, CSRM encounters challenges in accurately utilizing the neighbors’ information in the early stages of training, which can lead to erroneous results. 
It also fails to identify the most useful augmentation neighbors for the target user. Our  solves these problems by introducing a new vector burn-in technique for more precise cohort identification and a label-aware contrastive neighbor separation mechanism for mitigating the negative effect of noisy information. 
From Figure <ref>, we can see that  outperforms and is more consistent than both the previous neighbors-augmented learning methods in fraud detection.







 §.§ Influence of Cohort Size (RQ4)

In this section, we investigate the impact of cohort size variation on the effectiveness of . The results obtained from two fraud detection models on the Shopee-small dataset are illustrated in Figure <ref>. 
It is evident that the performance of the proposed label-aware cohort-augmented learning framework remains highly robust across varying cohort sizes once the cohort size exceeds a certain threshold. 
Specifically, AUC values demonstrate an increasing trend with cohort size, plateauing when K reaches four. With a smaller cohort, meaningful information is aggregated from training samples to enrich the representation of the target sample. 
However, as the cohort size increases, both valuable and noisy information is introduced to the target users, resulting in performance gains that saturate at K=4. 
Initially, most information learned from augmentation neighbors closely aligns with the target users, thereby enhancing their representation. 
Yet, with an increased cohort size, the likelihood of encountering augmentation neighbors with opposite labels to the target users rises, potentially compromising the effectiveness of newly augmented neighbors. 
However, the implementation of a label-aware neighbor separation scheme helps mitigate the adverse effects of noisy information introduced by the negative neighbors. 
When the impacts of negative neighbors and the label-aware neighbor separation scheme are comparable, the performance of  remains relatively stable across varying cohort sizes. In summary, our proposed  effectively enhances the learning capacity of base models and exhibits robust performance, irrespective of cohort size variations.



 §.§ Studies of Unveiling Camouflaged Fraud (RQ5)

To validate the 's effectiveness in unveiling the camouflage of frauds, we simulate fraudsters' camouflage strategies by selectively removing 3758 edges between certain fraudulent accounts and devices in the Shopee-graph dataset.
We then compare the performance of GraphSAGE and GeniePath with and without  integration on both the original Shopee-graph and the sub-sampled Shopee-graph datasets.
The comparison results in Table <ref> reveal that both GraphSAGE and GeniePath experience severe degradation when fraudster-device connections are deliberately removed.
However, the integration of  significantly enhances the effectiveness of GraphSAGE and GeniePath on the sub-sampled Shopee-graph. 
Specifically,  enables GraphSAGE to recover 86.3% (3242/3758) and GeniePath to recover 82.9% (3115/3758) of the deliberately removed connections.
This leads to R@P_0.9 improvements of 47.61% and 40.02%, highlighting ’s strength in identifying more fraudulent users and thus unveiling the camouflaged fraud.
Furthermore, the comparison between the base GNN models and their integration with  demonstrates that  can also augment the learning capabilities of GNN-based models in fraud detection.






 §.§ Efficiency of  with Integrating Vector Database

We assess the training and inference efficiency of  with and without Milvus in this subsection, integrating with MINT on the Shopee-large dataset. 

We performed experiments over 20 epochs for both cases and measured the throughput (number of training/inference tuples processed per second) for both CPU and GPU. 
To ensure fair comparisons, we maintained a consistent batch size of 512 across all experiments to maximize the GPU computational capacity. 
For Milvus, we employed HNSW as the index for vector search, known for its high accuracy and efficient indexing capabilities. 
In contrast, for the scenario without a vector database, we utilized PyTorch's built-in functions for cohort neighbor searches.
The throughput comparison results, as depicted in Figure <ref>, demonstrate that  exhibits significantly improved efficiencies when integrated with a vector database compared to  operating without one.
Specifically, the solution of  with Milvus achieves significantly higher training and inference efficiency, with speedups of up to 21.3x and 26.6x of CPU and GPU respectively, and incurs no extra infra investment.
Our experimental results highlight 's seamless integration with existing fraud detection methods and its effective combination with vector databases, providing a more efficient approach to fraud detection.






§ RELATED WORK


Fraud detection tasks analyze users’ behavioral data and transaction history to identify fraudulent activities, such as fake reviews, brushing, fraud payments, etc. 
Existing fraud detection systems fall into two categories: graph-based and sequence-based. 
Graph-based methods <cit.> use predefined features like devices, IPs, and purchase patterns to connect users, but these features may not reflect behavioral similarity and are easy to manipulate by fraudsters  <cit.>.
In contrast, sequence-based methods use recurrent neural networks (RNNs) <cit.> or attention-based mechanisms <cit.> to capture temporal and sequential behavior patterns, which are more robust against manipulation.
For example, some studies utilize behavior trees or Tree-LSTM <cit.> to extract user intentions from sequences, while others propose hierarchical explainable networks <cit.> to predict payment risks using critical historical behavioral sequences.

More recent works <cit.> design a general user behavior modeling framework to pre-train users’ behavior data and learn the robust representation for fraud detection.
However, sequence-based methods often neglect user interactions, which are crucial for detecting over-consistent behavior patterns in real-world fraud scenarios.

To address this, we propose a learning-based cohort augmentation method, which constructs personalized cohorts for each user based on the relevance of their behavioral sequences. 
In this manner, we can capture the subtle and dynamic patterns of fraudulent behaviors, and enhance the representation of each user.








§ CONCLUSIONS


This paper introduces , a label-aware cohort-augmented learning framework designed to enhance the fraud detection model’s capability in revealing camouflaged frauds.
Unlike conventional neighbor augmentation approaches,
 seamlessly integrates with existing models, offering a versatile solution to enrich user representations.
 incorporates a task-specific vector burn-in technique, enabling it to select augmentation neighbors and negative neighbors automatically for each training user. 
To refine the aggregated cohort information,  employs a label-aware neighbor separation mechanism, effectively distancing the negative neighbor representations in the augmented space. 
It addresses the limitations of conventional fraud detection methods by using learned behavior similarity instead of identity features to build connections between users, effectively handling noise and outliers, and explicitly considering label information during augmentation.
Comprehensive experiments and evaluations have demonstrated 's effectiveness and advantages in real-world fraud detection tasks. 
In summary,  provides a robust approach to enhance the distinguishing ability of existing fraud detection models, providing a promising strategy for detecting camouflaged fraudsters.


This research is supported by Shopee Singapore Private Limited, the Singapore Economic Development Board, and the Singapore Ministry of Education Academic Research Fund Tier 3 under MOE’s official grant number MOE2017-T3-1-007.


The work of Meihui Zhang is supported by the National Natural Science Foundation of China (62072033).
















ACM-Reference-Format












§ APPENDICES



§ PSEUDO-CODE OF 


Algorithm <ref> provides a detailed pseudocode description of the vector burn-in technique employed by . This technique represents a fundamental aspect of 's approach, effectively separating the optimization of embedding vectors from their utilization in subsequent cohort augmentation. Additionally,  stores the label information of training samples, a crucial element that plays a significant role in the subsequent negative neighbor separation process, enhancing its overall effectiveness.

On the other hand, Algorithm <ref> outlines the procedure for cohort identification and subsequent cohort-augmented training within the  framework. 
To enhance clarity, certain steps, including contrastive loss calculation, and vector regularization, have been omitted for simplicity.
During this process, it's essential to note that the parameters of auxiliary encoder g_aux(·) and the output layer used in the vector burn-in phase will remain fixed. This deliberate decision ensures the preservation of task-specific relevance between the target sample and the neighboring vectors, thereby facilitating the augmentation and separation of the representations between neighbors and target samples.
Furthermore, by retaining well-trained neighborhood information, the augmented-training process becomes significantly more manageable in comparison to dynamic neighbor identification and augmentation methods. This design choice not only promotes computational efficiency but also contributes to the stability and effectiveness of  in enhancing fraud detection.






each sample u_i ∈𝒯_train
Perform a forward pass: 𝐞_i = g_aux(𝐱_i) 
Search top-K similar vectors from ℰ_pool for u_i,   obtain 𝐄_i^a: 𝐄_i^a = search(𝐞_i, ℰ_pool) 
Search top-K similar vectors from ℰ_pool for u_i with   different labels,   obtain 𝐄_i^n: 𝐄_i^n = search(𝐞_i, ℰ_pool, 𝐲_𝐞≠𝐲_𝐮) 










§ MORE EXPERIMENTAL RESULTS


We also conducted experiments on two more methods widely compared recently, i.e., IHGAT <cit.> and FATA-Trans <cit.> to verify the effectiveness of .
The results in Table <ref> show that  considerably improves the detection performance of these two representative methods across all three datasets, with AUC (1% increase is considered significant <cit.>) 
and R@P_0.9 improvements up to 1.53% and 6.83%, respectively.
These new results further validate the effectiveness of our cohort augmentation for fraud detection.







§ VISUALIZATION OF 


We use t-SNE to visualize the user embeddings of a base model and its augmented version by , as shown in Figure <ref>.
We apply the same t-SNE settings for both cases, with a Euclidean distance metric and a perplexity of 30. 


Unlike most existing methods that use contrastive loss on the encoder outputs, our proposed  applies contrastive loss on the predicted logits. As a result, the cohort augmentation has a stronger effect on the first layer output of the output layer, which we use for visualization.

From Figure <ref>, we observe that fraudulent users do not always form a single cluster, but rather are dispersed across multiple groups of varying sizes.
This shows the variety and complexity of fraud schemes and user behaviors.
Normal users also have different types of clusters, indicating their different preferences and habits.



We also notice that the augmented model forms more distinct and compact clusters than the original base model.
This is because the personalized cohort augmentation can capture the similarities and differences among users, and create more homogeneous and isolated small clusters.
Furthermore, the augmented model can better distinguish fraudulent cohorts from normal user groups, owing to both the enhanced clustering precision conferred by cohort augmentation and negative neighbor separation techniques.

These techniques can push away similar but negative neighbors during the cohort-augmented training process, and make each cluster more cohesive and discriminative.






§ HYPER-PARAMETERS SETTINGS ON 


The hyperparameters α and β play a crucial role in 's objective function.
Specifically, they correspond to the weights of the supervised contrastive loss and auxiliary reconstruction loss, respectively. 
To provide a comprehensive understanding of how different settings of these hyperparameters affect the performance of , we conducted a detailed analysis.
This analysis focused on evaluating the changes in model performance when varying the values of α and β.
The experiments utilize base models HEN <cit.> and MINT <cit.> with Shopee-small dataset.
The results are summarized in Table <ref> and Table <ref>, and the analysis is presented as follows:







  * Study of α: An increase in α from its optimal value of 0.001 to 0.01 results in a marginal performance decline for both HEN with  and MINT with . 
A further increase of α to 1.0 leads to a more pronounced drop, due to an overemphasis on distinguishing training samples from negative neighbors, which could cause overfitting. 
On the other hand, a decrease in α below 0.001 reduces the model’s ability to counteract noise from negative neighbors, which will negatively impact ’s effectiveness.

  * Study of β: Adjusting the hyperparameter β also alters ’s performance.
A higher β diminishes the effectiveness of ’s cohort augmentation by aligning the main encoder too closely with the auxiliary encoder.
This will limit the diversity and robustness that cohort augmentation is intended to introduce.
Conversely, a reduction in β leads to a slight decrease in effectiveness, which is mitigated by the use of burn-in vectors from augmentation neighbors to enhance target sample representations.
These burn-in vectors help to enhance the representations of the target samples by introducing additional contextual information, which compensates to some extent for the decreased influence of the auxiliary reconstruction loss.


Overall, the performance of  remains relatively stable across different hyperparameter values. An α value of 0.001 and β 0.00001 achieve the best performance across different models. Hence, the objective function can be easily optimized without laborious hyperparameter tuning.



§ COMPUTATIONAL COMPLEXITY ANALYSIS OF 


The computational complexity of  involves both time and space complexity.
We denote the number of parameters P, the number of stored sample embedding vectors as n, the dimensionality of the embeddings as d, and the number of neighbors as k.




  * Time Complexity: The vector burn-in step’s time complexity is determined by the base model, with a complexity of O(T).
The cohort-augmented learning step’s time complexity is mainly dominated by the augmented learning model and the attentive neighbor aggregation layer, with a complexity of O(T) and O(k^2*d), respectively.
It also includes the search time for k augmentation neighbors using the HNSW index, with a complexity of O(k*log(n)).
So the time complexity should be O(T + k^2*d + k*log(n)).

  * Space Complexity: The space complexity in the vector burn-in step is O(P).
The cohort-augmented learning step involves storing burn-in vectors and constructing a graph-based index, leading to a complexity of O(nlog(n) + nd + P). 

Therefore, the overall space complexity is O(nlog(n) + nd + P).


Given that k is set to 5, the embedding dimension d is 64, and log(n) is usually small, the time complexity is mainly determined by the model rather than cohort neighbor searching.
The space complexity is dominated by the stored vector size and the model size, which are typically moderate in our anti-fraud system.




 Paper Title: Tamper-Resistant Safeguards for Open-Weight LLMs.



[NO \title GIVEN]
    [NO \author GIVEN]
    August 5, 2024
======================




[1]Equal contribution. ^Equal advising. Correspondence to mailto:rishubt2@illinois.edurishubt2@illinois.edu.
[2]Part of work done during an internship at the Center for AI Safety.





§ ABSTRACT 
Rapid advances in the capabilities of large language models (LLMs) have raised widespread concerns regarding their potential for malicious use. Open-weight LLMs present unique challenges, as existing safeguards lack robustness to tampering attacks that modify model weights. For example, recent works have demonstrated that refusal and unlearning safeguards can be trivially removed with a few steps of fine-tuning. These vulnerabilities necessitate new approaches for enabling the safe release of open-weight LLMs. We develop a method, called , for building tamper-resistant safeguards into open-weight LLMs such that adversaries cannot remove the safeguards even after thousands of steps of fine-tuning. In extensive evaluations and red teaming analyses, we find that our method greatly improves tamper-resistance while preserving benign capabilities. Our results demonstrate that tamper-resistance is a tractable problem, opening up a promising new avenue to improve the safety and security of open-weight LLMs.






§ INTRODUCTION







The most capable open-weight large language models (LLMs) released over the past year now rival closed-source frontier models <cit.>. The availability of open-weight LLMs for anyone to download and use has yielded numerous benefits, including lowering costs for end users and enabling academic research on safety and security <cit.>. However, as these models become increasingly powerful, many have raised concerns that they could be repurposed by malicious actors to cause harm, motivating research on how to safeguard these models against malicious use.


Existing open-weight models often adapt safeguards designed for closed-weight models served through APIs <cit.>. These safeguards include refusal mechanisms and preference-based training, and they have provided substantial robustness against input-based jailbreaking attacks. However, recent work has demonstrated these safeguards are trivially defeated by attacks that edit model weights, breaking down after only a handful of fine-tuning steps <cit.>. This poses a serious problem for open-weight models, because adversaries have full access to model weights and can tamper with
built-in safeguards.

The vulnerability of open-weight models to tampering attacks poses risks for model developers as well. Under background tort law, AI developers must exercise reasonable care, meaning they have an obligation to take reasonable precautions to prevent foreseeable harm. If malicious actors can easily customize models to cause critical harm, model developers may inadvertently violate reasonable care standards and become open to liability under existing law. Thus, there is an urgent need for more robust safeguarding techniques that can withstand tampering attacks.

In this work, we study the problem of tamper-resistant safeguards for LLMs. This problem is depicted in Figure <ref>. Unlike existing research on LLM safeguards, we focus on attacks that modify model weights, which we refer to as tampering attacks. This problem has been considered very challenging and by some intractable, as no method has yet provided substantial robustness to these attacks. However, making progress on this problem would provide a valuable tool to regulators and model developers by ameliorating the dual-use dilemma of open-weight models <cit.>. 

To demonstrate that progress on this problem is possible, we develop the first LLM safeguards that obtain strong robustness against a wide variety of tampering attacks. Our approach allows developers to add a safeguard such that tampering attacks cannot easily remove the safeguard, while preserving the general capabilities of the LLM. We achieve this by performing adversarial training against tampering attacks, leveraging approaches from meta-learning. We identify various crucial factors that enable our method to work, including the choice of tamper-resistance loss, the selection of train-time adversaries, and the two-stage approach that we use for building in safeguards.



We apply our method to develop tamper-resistant unlearning and refusal safeguards. In experiments, we demonstrate that our safeguards are far more robust to tampering attacks than prior methods. We stress-test our safeguards with extensive red teaming evaluations against  test-time adversaries, demonstrating resistance to fine-tuning attacks up to 5,000 steps. We hope our results foster future work on this important problem. Our experiment code and models are available at <https://github.com/rishub-tamirisa/tamper-resistance>.





§ RELATED WORK



  
Adversarial attacks on LLMs. Due to the extensive pre-training distribution of modern LLMs, they are prone to generating harmful content <cit.>. To mitigate this, many LLMs undergo fine-tuning to implement safeguards <cit.>, using methods such as reinforcement learning from human feedback (RLHF) <cit.> and direct preference optimization (DPO) <cit.>. While effective for normal use, these safeguards have been shown to be brittle, breaking down under jailbreak attacks <cit.> or a handful of fine-tuning steps on “uncensored” data <cit.>. This suggests current techniques for LLM alignment are inadequate, raising security concerns after deployment.

r0.5
    

    
    < g r a p h i c s >

    Comparison of our method to 12 baseline safeguards. Unlike prior methods, provides far greater tamper-resistance at similar levels of general capability, measured via MMLU. Tamper-resistance is computed as the normalized error on WMDP Biosecurity, Chemical Security, and Cybersecurity questions <cit.>, averaged across up to  fine-tuning attacks.
    
    







  
LLM safeguards. Since the discovery of these attacks, many safeguards have been proposed to defend against them. Against jailbreak attacks, defenses include system-level defenses <cit.> that modify or filter model inputs or outputs and model-level defenses such as adversarial training <cit.>. Alternatively, some works explore machine unlearning as a way to remove harmful knowledge entirely with techniques such as influence functions <cit.>, maximizing loss on forget sets <cit.>, or modifying representations <cit.>. However, jailbreaking defenses are not fully robust to adaptive adversaries <cit.>, and existing unlearning methods are not robust to adversaries with access to model weights <cit.>.





  
Robust safeguards. Several works have explored the tamper-resistance of unlearning methods for image classification <cit.>. For bidirectional BERT-style models, <cit.> proposed a meta-learning approach for robustly preventing models from learning harmful tasks. In concurrent work, <cit.> proposed a method extending this approach to small-scale vision classifiers and diffusion models. Recently, <cit.> discussed the potential for robust unlearning in LLMs to improve the safety of open-source models, and <cit.> proposed evaluation metrics for robust unlearning in LLMs. To the best of our knowledge, no methods have been proposed for autoregressive LLMs that are robust to tampering attacks. 

Several concurrent works have explored ways of defending LLM refusal mechanisms against fine-tuning. <cit.> add a perturbation loss to keep new embeddings close to the original, while <cit.> regularize harmful representations to look random. Unfortunately, the authors acknowledge that these approaches have highly limited robustness. We corroborate this in our comparisons, finding that they lack robustness to the tampering attacks in our evaluations.


























§ TAMPER-RESISTANT SAFEGUARDS






 §.§ Threat Model


We assume the defender releases an LLM with weights θ_G and a safeguard G applied. The defender's goal is to design G such that θ_G obtains high values on (θ_G) and (θ_G). Moreover, the defender seeks to preserve a high value of (θ_G) after the adversary's move. We consider a compute-bounded adversary with unrestricted access to θ_G, enabling attacks that directly modify θ_G. We refer to these as “tampering attacks.” The adversary's goal is to obtain a model θ'_G that minimizes the safety metric given reasonable compute limits, such as fine-tuning for 1,000 to 5,000 steps. We assume the adversary will not spend a significant fraction of the compute required to pre-train the LLM, since at that point they could train their own model without safeguards.









 §.§ Problem Definition and Metrics



We describe a general notation for quantifying the tamper-resistance of safeguards. Define G, θ_G, , and  as in the threat model. Let  denote a compute-bounded adversarial attack that maps θ_G to θ'_G, with stronger attacks obtaining lower values of (θ'_G). We say that a safeguard G is tamper-resistant if its post-attack (θ_G') is high across a broad range of strong test-time adversarial attacks 𝒜_test.

Note that θ_G often modifies an underlying θ that lacks safeguards, often through a fine-tuning procedure. Additionally, strong tamper-resistance can be obtained if the safeguard simply overwrites θ with noise, but this model would no longer be useful. Thus, maintaining a high (θ_G) is crucial, and evaluation of a safeguard must consider both its tamper-resistance and how well it preserves general capabilities. 


We focus on two common safeguard domains: weaponization knowledge restriction and harmful request refusal. In each domain, we define safety and capabilities test metrics, which we use alongside test-time adversaries to evaluate tamper-resistant safeguards.



  
Weaponization knowledge restriction. In weaponization knowledge restriction, safeguards prevent the model from producing text about weaponization knowledge, while preserving capabilities for benign knowledge domains. Existing safeguards of this nature include representation engineering methods like circuit breaking <cit.>. The  is defined as error on a forget set, and the  is defined as accuracy on a retain set. Specifically, we consider the problem of restricting biosecurity, chemical security, and cybersecurity knowledge, and evaluate the resulting model on the Weapons of Mass Destruction Proxy (WMDP) benchmark <cit.>. WMDP contains 3,668 multiple-choice questions, spanning biosecurity, chemical security, and cybersecurity knowledge. Importantly, WMDP questions do not evaluate hazardous knowledge directly, but instead measure proxy expert-level knowledge for each hazardous domain, such that suppressing the expert-level knowledge would also suppress the hazardous knowledge. We define the forget set as the respective hazardous knowledge subject in WMDP, and retain set as the complement of the given subject in MMLU <cit.>, a multi-task question-answering benchmark spanning 57 tasks across a variety of knowledge domains.



  
Harmful request refusal. In the harmful request refusal setting, safeguards prevent the model from producing “harmful” outputs. We define the  as the complement of average Attack Success Rate (ASR) of various jailbreaking attacks, while the  captures the conversational abilities of θ_G. Specifically, we use a static set of test cases from HarmBench, an automated red-teaming framework for measuring prompt jailbreak robustness in LLMs, to evaluate jailbreak ASR <cit.> after tampering attacks. We use MT-Bench, a multi-turn question-answering benchmark graded by an LLM judge, to evaluate conversational abilities <cit.>.


    


 §.§ Red Teaming




To properly measure the robustness of tamper-resistant safeguards, we conduct red-teaming with up to  adversaries, including many that are unseen at training time. In our evaluations, we subject our method to adversaries with varying compute budgets, access to held-out datasets, and diverse hyperparameters. For fine-tuning adversaries, we vary the learning rate, learning rate scheduler, optimization algorithm, and batch size. Many of these adversaries were fixed during early experiments, with some added over time as we found attacks that broke intermediate versions of our method. Extensive stress testing of this nature is critical for obtaining confidence in a tamper-resistant safeguard. For research on developing these safeguards, extensive red teaming also allows measuring incremental progress, using the number and strength of existing attacks one can defend against as a robustness metric.












§ SAFEGUARD TAMPER-RESISTANCE TRAINING









To obtain tamper-resistant safeguards, we propose a new method outlined in Algorithm <ref> inspired by adversarial training and meta-learning to directly strengthen LLM safeguards against tampering attacks, called Tampering Attack Resistance (). We identify unique properties of this adversarial training regime and leverage them to improve robustness.

Our method for training tamper-resistant safeguards consists of two phases: (1) model safeguarding and (2) tamper-resistance training.



 §.§ Model Safeguarding
 The method begins by including an initial safeguard G into a base model θ. For example, initial safeguards for knowledge restriction can be drawn from a wide variety of existing methods, including circuit breaking <cit.> or constrained gradient ascent for a particular knowledge domain. Similarly, we can include a refusal safeguard by performing RLHF <cit.> or DPO <cit.> on refusal completions. Importantly, these initial safeguards do not need to be tamper-resistant. Empirically, we find that this safeguarding step is important for obtaining strong tamper-resistance.







 §.§ Tamper-Resistance Training






Starting from θ_G_0, we train the tamper-resistant θ_G using a novel adversarial training procedure. Namely, we train against a set of tampering attacks 𝒜_train, where the defender's objective is to maximize a proxy  after applying an adversarial attack ∼𝒜_train to θ. Since it may not be feasible to differentiate through , we draw on insights from prior work in meta-learning, defining (θ_G) = θ'_G = θ_G + '(θ_G) as a perturbation on top of initial parameters, where backpropagation through ' is approximated with a straight-through estimator <cit.>. 








We focus on supervised fine-tuning (SFT) adversaries where  applies several steps of optimization to θ_G, which allows straight-through estimation through ' to benefit from the setting and approximations of first-order MAML <cit.>. However, we note key differences in our approach from standard meta-learning and prior methods <cit.>. In particular, traditional meta-learning techniques seek to obtain a model initialization that is close to optimality on multiple test distributions. In our setting, we seek to obtain an initialization that is far from optimality on multiple adversaries' test distributions. Novel to our approach in this new setting is the use of a tamper-resistance loss in the “outer loop” that differs from the fine-tuning adversary's loss function and serves to maximize the proxy safety metric. We depict this structure in Algorithm <ref>, and explain the objective below.





  
Impeding the adversary's loss. The aim of tamper-resistance training is to prevent adversaries with large compute budgets from reducing the  at test-time. In adversarial training for tamper-resistance, we define a tamper-resistance loss ℒ_TR that counters . We operationalize our goal of avoiding adversary optimality as searching for θ such that ℒ_TR is minimized for (θ).

Empirically, we find that the choice of tamper-resistance loss ℒ_TR significantly affects this goal. Prior work <cit.> negates the loss of a fine-tuning adversary, in which the aim is to arbitrarily maximize the adversary's loss throughout fine-tuning. This formulation has two issues: (1) maximizing a cross-entropy loss can cause divergence; (2) empirically we observe that when using this objective against fine-tuning adversaries, the model learns to explode the adversary's loss for the first few inner loop steps, while loss at later steps remains low. In Figure <ref>, we show the difference in choosing ℒ_TR to be a clamped negative cross-entropy loss vs. negative entropy loss for weaponization knowledge restriction. For the latter, ℒ_TR is eventually satisfied for all inner loop steps. For harmful request refusal, we choose ℒ_TR to be the DPO loss <cit.>. We provide further detail on the choice of ℒ_TR in both settings in Appendix <ref>.



  
Tamper-resistance objective. We now describe the general proxy objective used for preventing a tampering attack from recovering weaponization knowledge or harmful behavior. For a given , let 𝒟_TR and ℒ_TR respectively be a dataset and loss function such that minimizing ℒ_TR(θ_G, 𝒟_TR) serves as a proxy objective for maximizing the (θ_G). We define 𝒟_retain and ℒ_retain correspondingly for (θ_G). The defender's objective is to solve the following optimization problem:

    min_θλ_TR·𝔼_∼𝒜_train[  ℒ_TR((θ); 𝒟_TR) ] + λ_retain·ℒ_retain(θ; 𝒟_retain),
 where ℒ_TR is a tamper-resistance loss that counters (θ). The ℒ_retain term is a representation engineering <cit.> inspired retain loss for preserving performance on the capabilities proxy dataset 𝒟_retain, given by

    ℒ_retain(θ; 𝒟_retain) = 𝔼_x∼𝒟_retain[ℒ_LM(θ, x) + ‖ h_θ(x) - h_θ_G_0
    (x) ‖_2^2]
 where h_θ(·) returns the residual stream hidden states for model parameters θ and ℒ_LM is the standard language modelling cross-entropy loss. Empirically, we find that pushing retain-set residual stream representations to be close to the base model θ_G_0 via the ℓ_2-norm loss in Equation <ref> maintains a high (θ_G). In Equation <ref> we include λ_TR and λ_retain as scalar weightings for the tamper-resistance loss and retain loss, respectively. We provide further details on the design of the tamper-resistance loss function in Appendix <ref> as well as an efficiency trick for sampling fine-tuning attacks for in Appendix <ref>.



























































§ EXPERIMENTS


We evaluate in weaponization knowledge restriction and harmful request refusal settings, with results shown in Table <ref> and Table <ref> respectively. We discuss the setup, baselines, and analysis for our results. In each setting, we use a specific set of training adversaries 𝒜_train and test adversaries 𝒜_test. Further experiment details are presented in Appendix <ref>.



 §.§ Weaponization Knowledge Restriction

We now describe the setup, baselines, and results for our weaponization knowledge restriction experiments, including the knowledge domains, optimizers, and evaluation details.



  
Setup.
 We focus on implementing tamper-resistant safeguards for restricting proxy weaponization knowledge about biosecurity, chemical security, and cybersecurity from Llama-3-8B-Instruct <cit.> that has been initially safeguarded via the Random Mapping method discussed in Appendix <ref>. For each weaponization domain, we assign 𝒟_𝒯ℛ to the corresponding forget set described in Appendix <ref>. We proceed to sample train-time 64-step fine-tuning attacks from different data distributions, detailed in Appendix <ref>. We use N = 750 outer loop steps, ScheduleFree AdamW <cit.> with a learning rate of 2×10^-5 as the outer loop tamper-resistance optimizer, and loss scales of λ_TR = 4.0, λ_retain = 1.0. Lastly, We evaluate Pre-Attack and Post-Attack accuracy on corresponding WMDP subjects <cit.> averaged across all adversaries in Appendix <ref>, and measure benign capabilities via the complement of subjects related to each proxy weaponization domain in MMLU <cit.>.







  
Baselines.
We evaluate two recently proposed knowledge restriction methods: RMU <cit.> and LLMU <cit.>. We also design two baseline methods for knowledge restriction: Min Posterior, which minimizes posterior loss on forget set tokens; Max Entropy, which maximizes entropy on forget set tokens. Two additional methods, MLAC <cit.> and SOPHON <cit.>, require substantial modifications for the LLM setting, so we show results on adapted versions of these baselines in <Ref>.









  
Results.
We show weaponization knowledge restriction safeguard results on Llama-3-8B-Instruct in <Ref> and Figure <ref>. These results are averaged across all 28 adversaries from Appendix <ref>. Our large-scale experiments corroborate the findings in recent work that existing LLM safeguards are extremely brittle to fine-tuning attacks. By contrast, maintains low post-attack forget accuracy across all three domains. We also observe that maintains non-negligible retain accuracy, although there is a modest trade-off, much like how adversarial robustness trades off benign accuracy. In Figure <ref>, we observe that is robust to significantly more fine-tuning attacks than all prior methods. While existing baselines break down under most attacks, obtains a post-attack forget accuracy near random chance for nearly all attacks, indicating a successful defense.

Overall, we find that provides significantly more robustness to realistic fine-tuning attacks than all prior methods, including attacks as strong as 5,000 steps of optimization on completely held-out data. These results demonstrate for the first time that obtaining strong tamper-resistance for open-weight LLMs may be possible.





 §.§ Harmful Request Refusal







 We now describe the setup, baselines, and results for our harmful request refusal experiments, including the datasets used and evaluation details. 


  
Setup.For harmful request refusal training, we seek to make existing refusal safeguards in Llama-3-8B-Instruct robust to tampering attacks. We sample train-time adversaries that perform 64-step SFT attacks using the Anthropic-HH-RLHF dataset <cit.>, following the methodology in Appendix <ref>. Similar to the weaponization knowledge restriction setting, we use N = 100 outer loop steps, ScheduleFree AdamW <cit.> with an LR of 6 × 10^-5 as the outer loop tamper-resistance optimizer, and loss scales of λ_TR = 0.1, λ_retain = 1.0. We evaluate the Post-Attack jailbreak attack success rate (ASR) on HarmBench <cit.> after the tampering attacks in Appendix <ref>, and measure benign capabilities preservation via MT-Bench <cit.>, which evaluates multi-turn conversation ability. 






  
Baselines. We consider 4 baselines alongside our model: Llama-3-8B-Instruct (Refusal Trained); Representation Rerouting (RR) <cit.> on Llama-3-8B-Instruct, which trains to push representations for harmful input prompts to be orthogonal to the original corresponding representations in Llama-3-8B-Instruct; R2D2 <cit.> on Zephyr-7B <cit.>, which performs adversarial training against GCG attacks <cit.>; and RepNoise <cit.> on Llama-2-7B <cit.>, which regularizes harmful representations to Gaussian noise.  





  
Results.
We show refusal results in <Ref>. While the Refusal Training, RR, and R2D2 baselines resist jailbreak attacks in HarmBench before tampering, we find that percentage attack success rate jumps up to above 77 after tampering, while our method only rises to 61.7. Since we apply our refusal safeguard to Llama-3-8B, it does reduce MT-Bench by 1.7. However, this exceeds the MT-Bench score of fairly capable open-weight models, indicating that benign capabilities are largely preserved. We leave the exploration of the full impact on capabilities to future work. Additional results are in <Ref>. In general, we find that our model refuses more Post-Attack jailbreaks than previous baselines, and demonstrates the flexibility of the tamper-resistance objective to accommodate the harmful request refusal setting.









 §.§ Analysis





  
Red teaming.
To stress-test the tamper-resistance of our models, we conduct an extensive suite of supervised fine-tuning attacks with  distinct adversaries. We vary the optimizer, number of optimization steps, learning rate, learning rate schedule, fine-tuning dataset, batch size, and overall fine-tuning method (e.g., full fine-tuning versus parameter-efficient fine-tuning). By default, our attacks use 1,000 fine-tuning steps, although some use 5,000 steps. Full details for these adversaries are provided in <Ref>.

We show red teaming results in Figures <ref> and <ref>. While baseline safeguards withstand fine-tuning attacks in a small number of cases, most adversaries succeed in removing the safeguards. By contrast, our safeguard is robust to a wide range of adversaries. This shows that tamper-resistance is a tractable problem on which progress can be made. However, our method is currently not robust to parameter-efficient fine-tuning (PEFT) attacks (adversaries 27 and 28), highlighting the importance of extensive red teaming when developing tamper-resistant defenses. We hypothesize that future work could easily address this issue, as we demonstrate in <Ref> that targeted patching of vulnerabilities is possible.


r0.5
    

    
    < g r a p h i c s >

    Our safeguard remains robust to fine-tuning attacks that greatly exceed the 64 steps used by train-time adversaries. For the LLMU safeguard, the adversary's loss quickly decreases into the recovery region. By contrast, maintains flatness at a high loss for all 1,000 steps.
    

    





  
Generalization to stronger test-time attacks. In Figure <ref>, we show a fine-tuning attack at an LR of 2×10^-5 that attempts to recover biosecurity knowledge on our model and a model safeguarded with LLMU. We find that the tamper-resistance of generalizes far beyond the 64 steps used by train-time adversaries. Surprisingly, we observe that the test-time adversary's cross-entropy loss does not decrease below 7 for all 1,000 steps. Moreover, the loss enters a plateau and does not decrease at all after 200 steps. In <Ref>, we find that the length and height of this plateau can be increased by increasing the number of inner loop steps during adversarial training. As a point of reference, we show the progression of the same attack on LLMU. In this case, the adversary's loss decreases to within the recovery region in under 20 steps. We note that the adversary's test loss decreasing into the recovery region does not always correspond with recovery on downstream metrics (e.g., WMDP), but often does.































    



























§ CONCLUSION


We introduced a novel method for implementing tamper-resistant safeguards for LLMs and explored applications in weaponization knowledge restriction and harmful refusal training. We compare our results to prior work in each setting, finding that our method is the first method robust under the rigorous red-teaming evaluation that we consider. More broadly, we demonstrate that progress on open-weight tamper-resistance is tractable. We believe this line of research is crucial for enabling ongoing deployment of robust, open-weight LLMs, ensuring their alignment with regulatory frameworks and preemptively addressing the risk of malicious use.





  
Acknowledgements. We thank Steven Basart and Luis Fernandez for providing valuable feedback for the paper, and Andriy Novykov for providing support for compute resources at the Center for AI Safety. We thank the Center for AI Safety for providing significant compute resources for this project. We also thank Volodymyr Kindratenko and the National Center for Supercomputing Applications (NCSA) and Illinois Campus Cluster Program (ICCP) for supporting our computing needs. This work used NVIDIA GPUs at NCSA Delta through allocations CIS230117 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which is supported by NSF Grants #2138259,
#2138286, #2138307, #2137603, and #2138296.











abbrvnat










§ LIMITATIONS




Our method for training tamper-resistant safeguards demonstrates considerable robustness against a wide range of tampering attacks, yet several avenues for improvement remain: (1) While we focus on supervised fine-tuning attacks, the broader spectrum of open-weight tampering techniques necessitates diverse future red-teaming efforts. (2) Scaling to larger models poses computational challenges that require optimization to reduce overheads. 



Tamper-resistance alone cannot fully mitigate the risks of malicious AI use. While it raises the initial costs for adversaries, it can eventually be circumvented. Once open-weight models are released, they cannot be “unreleased,” leaving any compromised defenses permanently vulnerable. Therefore, tamper-resistance should be considered a supplement to the broader effort of of improving the offense-defense balance of AI systems. Addressing these limitations will improve the robustness of LLMs to tampering and better support open-weight model developers.







§ METHOD DETAILS





 §.§ Initial Weaponization Knowledge Restriction Safeguard


Prior to tamper-resistance training, we install a safeguard that achieves surgical knowledge restriction on the target hazardous domain. Let h_θ(𝒟) denote the distribution of post-decoder layer residual stream activations for input sequences sampled from some data distribution 𝒟 and model weights θ. We define  (x) for some input sequence x, which returns fixed Gaussian-sampled vectors that are chosen via hashing the corresponding input token for each residual stream index of x in θ. As a proxy for scrubbing target representations according to downstream task labels, we propose a weaponization knowledge restriction safeguard termed Random Mapping, which maps h_θ(𝒟_TR) to random noise as follows:



    min_θ𝔼_x∼𝒟_TR[1 - | h_θ(x) ·(x)/‖ h_θ(x)‖‖(x)‖| ] + ℒ_LM(θ; 𝒟_retain)


The objective of Equation <ref> maximizes cosine similarity between row vectors in the residual stream in every layer of the LLM from h(𝒟_TR) and the hashed random vectors from (·). By providing each token's residual stream a unique random vector to push toward, the loss encourages a “re-mapping” of token representations from 𝒟_TR to the noised vectors. We include an additional term for preserving performance on 𝒟_retain via the language-modelling cross-entropy loss ℒ_LM. We show the performance of the raw Random Mapping safeguard as an ablation in Table <ref>, listed as “Excl. Adv. Training.”




 §.§ Designing the Tamper-resistance Loss




  
Weaponization knowledge restriction. For weaponization knowledge restriction, we summarize our intuition for ideal tamper-resistance loss design, corroborated by our empirical findings in Figure <ref>: we seek to flatten the adversary's loss at a high value, rather than simply raise its y-intercept.

We choose the tamper-resistance loss ℒ_TR as an entropy loss to be maximized during the adversary's cross-entropy fine-tuning trajectory, since maximizing entropy would impede the adversary's cross entropy loss from decreasing during fine-tuning. In other words, we wish to obtain θ such that after an adversary performs a fine-tuning attack on θ via a cross-entropy loss, entropy is still high. We find that this formulation achieves the desired flattening behavior, and we depict the difference in flattening between the choosing ℒ_TR to be a negative cross-entropy loss and negative entropy loss in Figure <ref>. In the lefthand plot, where ℒ_TR is a cross-entropy loss, loss only increases in the first inner loop step. In the righthand plot, where ℒ_TR is a negative entropy loss, entropy is eventually maximized in all inner loop steps. Figure <ref> also demonstrates the generalization of the flat adversary loss behavior beyond the length of the simulated fine-tuning trajectories during .



  
Harmful request refusal. For harmful request refusal, we choose ℒ_TR to be the DPO loss <cit.>, which works as follows. Given a DPO dataset containing pairs of rejected and refusal completions, the sampled  performs SFT on rejected completions, and the tamper-resistance loss ℒ_TR is a DPO loss computed on the pair chosen and rejected completions on parameter coordinates along the  trajectory. This encourages to find an initialization θ such that after a harmful fine-tuning attack, the model still prefers refusal completions over harmful completions when given an harmful prompt. While this does not necessarily encourage a flat adversary loss, we find empirically in Table <ref> that this formulation increases the average (θ_G) defined in Section <ref> after fine-tuning attacks on harmful data, detailed in Section <ref>.



We also observe that the length of the simulated adversary SFT trajectory during training affects test-time generalization in both Figure <ref> and Appendix <ref>. In particular, larger values of K result in increased tamper-resistance for longer SFT attacks. However, to maintain reasonable runtime efficiency, we need a more efficient sampling technique than simply running K independent trajectories of varying length for every outer-loop step in Algorithm <ref>, which we describe in Section <ref>.



 §.§ Efficiently Sampling Fine-tuning Attacks
 


Optimizing Equation <ref> with gradient descent requires simulating K tampering attacks for each tamper-resistance optimizer update, which is prohibitively expensive to run when the sampled  performs SFT and θ contains billions of parameters. Inspired by prior work on snapshot ensembles <cit.>, we leverage an efficiency trick: we can reuse the coordinates along steps of a single adversary fine-tuning trajectory of length K to obtain K-1 additional (though non-independent) trajectories of increasing length. Using this trick, we collect all K parameter coordinates along the trajectory into a single batch for computing the tamper-resistance losses, effectively sampling  from 𝒜_train non-IID. To further improve runtime efficiency, we do not compute the tamper-resistance loss ℒ_TR on all K steps and instead sub-sample coordinates along the trajectory for computing ℒ_TR within an adversary batch, for example every 4 adversary optimization steps. Additionally, we reduce variance in the tamper-resistance gradient by computing the tamper-resistance loss at each inner loop step on the same held-out batch, denoted as x_TR in Algorithm <ref>. 



 §.§ Implementation Details and Resource Requirements


We perform training on Llama-3-8B-Instruct <cit.> with 8 NVIDIA 80GB A100 GPUs, leveraging distributed training via FSDP <cit.>. We use ZeRO Stage 3 from DeepSpeed <cit.>, which shards optimizer states, gradients, and parameters during training. While the efficiency trick in Appendix <ref> improves runtime, we note additional considerations for conserving GPU memory. 

First, simulating fine-tuning attacks that require additional state (e.g., momentum) in the inner loop of Algorithm <ref> requires initializing a fresh optimizer for every outer loop iteration. Since we use an outer-loop optimizer that also requires maintaining state (ScheduleFree AdamW <cit.>), we move the outer loop optimizer to the CPU before instantiating inner-loop optimizers.

Second, first-order meta-learning in smaller models can typically be implemented by running multiple forward passes for each inner loop iteration, averaging losses, then backpropagating on the averaged loss term. However, because each inner-loop tamper-resistance loss term (ℒ_TR in Algorithm <ref>) is computed on a separate forward pass, this requires maintaining K computation graphs in memory. Since this is infeasible on reasonable hardware for LLMs with billions of parameters, we circumvent this inefficiency by accumulating tamper-resistance gradients in a separate data structure (g_TR in Algorithm <ref>). We note that this can be done without using additional  and  distributed operations, since tamper-resistance gradient accumulation and application to the pre-inner loop model parameters (θ_i-1 in Algorithm <ref>) can be computed solely on sharded gradients. 












§ ADDITIONAL EXPERIMENTS





 §.§ Benign Fine-Tuning






An important property of open-weight models is that they can be fine-tuned to improve performance on custom data or in specific domains. Thus, ideal tamper-resistant safeguards should allow continued fine-tuning of a model while preserving the safeguard. We evaluate whether models can be fine-tuned on data unrelated to the safeguard using economics as an example domain. Using models with biosecurity and cybersecurity safeguards, we perform supervised fine-tuning on the WMDP auxiliary economics corpora <cit.>. We fine-tune models for 800 steps using a learning rate of 2×10^-6 and a batch size of 64. For the test set, we evaluate average accuracy on the corresponding MMLU High School Macroeconomics and Microeconomics subjects before and after fine-tuning. To confirm that the safeguard remains tamper-resistant in this setting, we also evaluate accuracy on corresponding WMDP subjects before and after fine-tuning.

We show the results of this evaluation in <Ref>. We find that accuracy on economics questions can be improved by 7.5 percentage points without recovering significant hazardous knowledge. This illustrates that strong tamper-resistance can be compatible with benign model editing. Additional experiment details are discussed in Appendix <ref>.



 §.§ Varying the Train-time Inner-loop Length K



r0.5
    

    
    < g r a p h i c s >

    Comparison of a 1,000-step SFT attack against with the inner-loop length K varied between {8, 16, 32, 64}. Test-time loss plateau magnitude and duration tend to increase as K increases.
    




Recall that via the efficiency trick discussed in Appendix <ref>, a single inner loop trajectory of length K during returns the K sampled attacks in Algorithm <ref>.  We compare the test-time loss robustness as we vary the length of the inner loop K during , running fine-tuning attacks for 1,000 steps on a held-out forget dataset for biosecurity weaponization (Adversary 8 in Table <ref>). For each value of K, we observe a plateau in the test loss that drops off at later steps as K increases. This suggests that the robustness of improves as the inner loop length increases. Prior work also corroborates that increasing the inner-loop length during meta-learning increases test-time generalization <cit.>. We note the contrast to conventional meta-learning methods mentioned in Section <ref>, in which typical meta-learning applications seek optimality after as few test-time steps as possible <cit.>. Here, our results suggest that the objective is incentivized to run with as many inner loop steps as possible. We find that K=64 provides significant robustness to the range of adversaries we consider in Section <ref> and Appendix <ref>, while balancing computational efficiency as discussed in Appendix <ref>. 






 §.§ Targeted Patching of Vulnerabilities



In Section <ref>, we discuss the set of adversaries sampled during for each weaponization knowledge domain. However, we included the “Retain-set SFT followed by Weaponization-domain SFT” (R→F in Table <ref>) adversary to improve robustness to specific adversaries that broke the defense in earlier iterations of the method that did not sample this adversary. We find in Table <ref> that including the R→F adversary significantly improves robustness to more test-time attacks, indicated by an improvement in Post-Attack Forget error by over 20.0 percentage points. Importantly, we find that including the 64-step R→F adversary at train-time improves test-time robustness at for the 1000-step R→F adversary by nearly 40.0 percentage points, indicating that target patching of vulnerabilities is possible within .










 §.§ Ablations






  
Retain loss ℓ_2 representation engineering term. We observe in Table <ref> that retain-set MMLU performance is worse when excluding representation engineering <cit.> retain loss from . When including the term, retain MMLU improves by at least 5.0 percentage points. 



  
Including the initial safeguard. In Table <ref>, we examine the impact of incorporating the Random Mapping safeguard step prior to the adversarial training phase during . The Random Mapping safeguard in isolation achieves a near-random chance Pre-Attack Forget accuracy of 27.3. However, it is susceptible to fine-tuning attacks similar to other baselines in Table <ref>, indicated by a higher Post-Attack Forget accuracy of 59.4. When including the tamper-resistance adversarial training phase, we observe significantly increased tamper-resistance as the Post-Attack Forget accuracy decreases by nearly 28 percentage points. 

We also examine the impact of excluding the initial safeguarding step (“Excl. Initial Safeguard” in Table <ref>), finding that Pre-Attack and Post-Attack Forget accuracies are both higher 
without the initial safeguard. While the tamper-resistance adversarial training phase achieves moderate tamper-resistance without the initial safeguard step, we find that including the Random Mapping phase improves downstream tamper-resistance by 15.6 percentage points.





  
Varying the tamper-resistance loss scale λ_TR. We compare the downstream robustness of when varying the tamper-resistance loss weighting λ_TR between 1.0 and 4.0 in Table <ref>. We observe that when setting λ_TR = 1.0, maintains high retain MMLU accuracy at 62.5 percentage points, with moderate tamper-resistance indicated by a Post-Attack Forget accuracy of 40.8. Further increasing λ_TR to 4.0 in our final model results in a significantly improved Post-Attack Forget Accuracy of 31.3, with a partial decrease in Retain MMLU to 54.9. When varying λ_TR, we keep λ_retain constant; thus, our results indicate a clear way to increase downstream tamper-resistance by increasing the weighting of the tamper-resistance gradient during , reflecting a balance between tamper-resistance and capabilities similar to the robustness-performance tradeoff for adverarial robustness in vision models.





§ EXPERIMENT DETAILS





 §.§ Weaponization Domain Proxy Dataset Details




  
Biosecurity. We use a synthetically labeled partition of the Pile <cit.> that filters for relevance to biology and the Camel AI Biology dataset <cit.>. We generate synthetic labels for Pile token sequences using openchat-3.5 <cit.>, categorizing them as "Cellular Biology" or not. This process yields 49,984 samples: 7,558 for the forget-set (Pile-bio Forget) and 42,426 for the retain set (Pile-bio Retain). Concurrently, we pack entries from the Camel AI Biology dataset to the truncation-enabled 256 tokenization limit, resulting in 54,258 samples of about 188 words each (Camel-bio Forget). We apply the same procedure to our held-out hazardous biology dataset (identical to the WMDP biosecurity Forget-set), producing 598,933 samples of similar length (OOD Forget).



  
Chemical Security. We use a private forget dataset containing text sequences about hazardous chemical security content (Chem Forget).



  
Cybersecurity. We scrape CTF writeups on CTFtime <cit.> that are numbered between 1 and 39181, collecting cybersecurity writeups written as recently as 2024. We filter to keep writeups that contain more than 150 characters. As a result of filtering and HTTP errors while scraping, our resulting forget dataset contains slightly over 18k samples (Cyber Forget).




 §.§ Train-time Settings and Adversaries








  
Weaponization knowledge restriction. For each weaponization knowledge restriction domain, we have a corresponding retain dataset 𝒟_retain, comprised of a mix of data from the Pile-bio Retain set and Magpie-Align instruction-tuning dataset <cit.>, which we refer to as the Retain-set.  Referencing the datasets described in Appendix <ref>, we specify the adversary tamper-resistance datasets that constitute the data used by attacks sampled from 𝒜_train during training, as well as the attack setup for each domain as follows:


    
  * Biosecurity: We simulate six adversaries from the following three setups: Pile-Bio Forget-set SFT, Camel-Bio Forget-set SFT, and Retain-set SFT followed by Pile-Bio Forget-set SFT (R→F), where the switching point between Retain-set and Pile-Bio Forget-set SFT within the 64-step length trajectory is sampled from a beta distribution β(6.0,3.0). LRs are sampled from {2 × 10^-5, 4 × 10^-5}.
    
  * Chemical Security: We simulate nine adversaries from the following three setups: Chem Forget-set SFT, Retain-set SFT, and Retain-set followed by Chem Forget-set SFT, using the same switching-point sampling scheme as in the Biosecurity setting. LRs are sampled from {2 × 10^-6, 2 × 10^-5, 4 × 10^-5}.
    
  * Cybersecurity: We simulate four adversaries from the following two setups: Cyber Forget-set SFT, and Retain-set SFT followed by Cyber Forget-set SFT, using the same switching point sampling scheme as in the Chemical Security and Biosecurity settings. LRs are sampled from {2 × 10^-5, 4 × 10^-5}.


For each weaponization knowledge domain, we create 80-20 splits for adversary and held-out data of the corresponding forget sets, respectively. For Biosecurity, which uses multiple forget datasets, this involves creating 80-20 splits for each dataset, then combining the corresponding splits. The adversary data splits are used for sampled attacks from 𝒜_train, whereas the held-out split is used for computing tamper-resistance losses. The held-out splits for each domain correspond to 𝒟_TR in Section <ref>. We use minibatches from a held-out dataset for computing tamper-resistance losses rather than cycling through a single dataset, following the recommendation of <cit.>, in which first-order meta-learning without properly held-out minibatches caused a performance degradation.

All train-time adversary setups are tabulated in Table <ref>, where F-Pile, F-Chem, F-Cyber denote the respective datasets described in Appendix <ref>, and Retain denotes the mixed Pile-bio and Magpie-Align Retain-set described in Appendix <ref>. We use R→F to label the adversaries that perform Retain-set SFT followed by Forget-set SFT. The final column is an abbreviation for Finetuning Paradigm and indicates whether the SFT setup used full parameter finetuning or parameter-efficient fine-tuning (PEFT) via LoRA adapters <cit.>.





@lccccccc@
Train-time adversary setups for weaponization knowledge restriction of Biosecurity, Chemical Security, and Cybersecurity.
 


Adversary     Dataset     Opt. Steps (K)     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8c
  – continued from previous page 

Adversary     Dataset     Opt. Steps     Opt.     LR     LR Schedule     Batch Size     FT Paradigm 



8rContinued on next page 
 




8cBiosecurity Weaponization Restriction 

Adv 1     F-Pile     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 2     F-Pile     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 3     F-Camel     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 4     F-Camel     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 5     R→F     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 6     R→F     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

8cChemical Security Weaponization Restriction 

Adv 1     F-Chem     64     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 2     F-Chem     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 3     F-Chem     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 4     Retain     64     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 5     Retain     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 6     Retain     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 7     R→F     64     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 8     R→F     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 9     R→F     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

8cCybersecurity Weaponization Restriction 

Adv 1     F-Cyber     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 2     F-Cyber     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 3     R→F     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 4     R→F     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 









  
Harmful request refusal. For harmful request refusal, we choose the retain dataset 𝒟_retain to be the Magpie-Align instruction-tuning dataset <cit.>. We sample train-time adversaries that perform K=64 steps of SFT on rejected completions of the Anthropic-HH-RLHF dataset <cit.> and vary the learning rate within {2 × 10^-6, 2 × 10^-5, 4 × 10^-5}. We depict this list of adversaries in Table <ref>. 

While the rejected completions from Anthropic-HH-RLHF constitute the data used for sampled attacks from 𝒜_train for harmful request refusal, we compute the tamper-resistance loss ℒ_TR as follows. Since ℒ_TR is the DPO loss <cit.> in this setting, we use the base model weights θ as the reference model and sample harmful and benign completions from a modified test split of Anthropic-HH-RLHF <cit.>, where rejected completions are replaced with refusals <cit.>. To avoid keeping the base model weights in memory and speed up training, we precompute the reference model DPO log-probabilities for the full Anthropic-HH-RLHF dataset before training. To summarize, we have that the sampled train-time adversaries perform SFT on the rejected completions from Anthropic-HH-RLHF, and tamper-resistance DPO loss is computed on the corresponding modified refusal completions; the modified refusal completions in this setting correspond to 𝒟_TR in Section <ref>.

In practice, we perform an additional 100 steps of supervised fine-tuning on the Magpie-Align dataset to improve the benign capabilities performance of the refusal model in Table <ref>.





@lccccccc@
Train-time adversary red-teaming setups harmful request refusal. “A-HH-Rejected” in the Dataset column corresponds to the adversary dataset comprised of rejected completions from the Anthropic-HH-RLHF dataset.
 


Adversary     Dataset     Opt. Steps (K)     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8c
  – continued from previous page 

Adversary     Dataset     Opt. Steps (K)     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8rContinued on next page 
 



8cHarmful Request Refusal 

Adv 1     A-HH-Rejected     64     AdamW     2×10^-5     No Warmup      64     Full Parameter 

Adv 2     A-HH-Rejected     64     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 3     A-HH-Rejected     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 













§ RED TEAMING DETAILS






 §.§ Weaponization Knowledge Restriction


In Table <ref>, we list all test-time adversary setups for recovering Biosecurity, Chemical Security, and Cybersecurity Weaponization knowledge. 

For Biosecurity, we examine post-attack forget accuracy after fine-tuning for 1,000 steps on three data distributions: the Pile-bio Forget set and the Retain-set used during Random Mapping and , and an OOD-Forget set mentioned in Appendix <ref>, which is unseen during tamper-resistance training. We conduct an additional attack using 5,000 optimization steps on this unseen forget distribution to compose our strongest adversary. We leverage a consistent setup for Chemical Security and Cybersecurity, substituting the Pile-bio Forget set for the respective Chem and Cyber Forget sets. We exclude the unseen forget distribution adversaries for a total of 24 attacks, as we were unable to obtain an equally high quality held-out dataset for Chemical Security and Cybersecurity.

For all subjects, we compare the effect of different optimizers (AdamW, Adadelta, and Stochastic Gradient Descent with Nesterov Momentum, Schedule Free AdamW) <cit.>, learning rates (2 × 10^-6, 2 × 10^-5, 4 × 10^-5), and learning rate scheduling techniques (SGDR and 30 steps of linear warmup) <cit.>. 

Similar to Table <ref>, we use F-Pile, F-Chem, F-Cyber, OOD-F, and Retain in the Dataset Column to denote the respective datasets described in Appendices  <ref> and <ref>. Separate from the R→F train-time adversary, R→F at test-time represents the combination of 0.4 × Opt. Steps of Forget-set SFT followed by 0.6 × Opt. Steps of Retain-set SFT. In practice, we found that this fraction composed a potent attack which we sought robustness against as described in Appendix <ref>.

In the Optimizer column, Schedule Free is an abbreviation of AdamW Schedule Free and SGD Nesterov is an abbreviation of SGD with Nesterov Momentum. In cases where the adversary used parameter-efficient fine-tuning (PEFT) via LoRA adapters <cit.>, we used a LoRA config with an attention dimension, or rank, of 16, a LoRA alpha value of 32, a LoRA dropout of 0.05, on target linear modules:      

{, , , ,, , }.













@lccccccc@
Test-time adversary red-teaming setups for weaponization knowledge restriction of Biosecurity, Chemical Security, and Cybersecurity.
 


Adversary     Dataset     Opt. Steps (K)     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8c
  – continued from previous page 

Adversary     Dataset     Opt. Steps (K)     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8rContinued on next page 
 



8cBiosecurity Weaponization Restriction 

Adv 1     OOD-F     5000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 2     OOD-F     5000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 3     F-Pile     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 4     F-Pile     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 5     Retain     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 6     Retain     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 7     OOD-F     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 8     OOD-F     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 9     R→F     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 10     R→F     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 11     F-Pile     1000     Adadelta     2×10^-5     No Warmup     64     Full Parameter 

Adv 12     F-Pile     1000     Adadelta     4×10^-5     No Warmup     64     Full Parameter 

Adv 13     F-Pile     1000     Schedule Free     2×10^-5     No Warmup     64     Full Parameter 

Adv 14     F-Pile     1000     Schedule Free     4×10^-5     No Warmup     64     Full Parameter 

Adv 15     F-Pile     1000     SGD Nesterov     2×10^-5     No Warmup     64     Full Parameter 

Adv 16     F-Pile     1000     SGD Nesterov     4×10^-5     No Warmup     64     Full Parameter 

Adv 17     F-Pile     1000     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 18     F-Pile     1000     AdamW     2×10^-6     30 Steps Warmup     64     Full Parameter 

Adv 19     F-Pile     1000     AdamW     2×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 20     F-Pile     1000     AdamW     4×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 21     F-Pile     1000     AdamW     2×10^-5     SGDR     64     Full Parameter 

Adv 22     F-Pile     1000     AdamW     4×10^-5     SGDR     64     Full Parameter 

Adv 23     F-Pile     1000     AdamW     2×10^-5     No Warmup     32     Full Parameter 

Adv 24     F-Pile     1000     AdamW     4×10^-5     No Warmup     32     Full Parameter 

Adv 25     F-Pile     1000     AdamW     2×10^-5     No Warmup     128     Full Parameter 

Adv 26     F-Pile     1000     AdamW     4×10^-5     No Warmup     128     Full Parameter 

Adv 27     F-Pile     1000     AdamW     2×10^-5     No Warmup     64     PEFT 

Adv 28     F-Pile     1000     AdamW     4×10^-5     No Warmup     64     PEFT 

8cChemical Security Weaponization 

Adv 1     F-Chem     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 2     F-Chem     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 3     Retain     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 4     Retain     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 5     R→F     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 6     R→F     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 7     F-Chem     1000     Adadelta     2×10^-5     No Warmup     64     Full Parameter 

Adv 8     F-Chem     1000     Adadelta     4×10^-5     No Warmup     64     Full Parameter 

Adv 9     F-Chem     1000     ScheduleFree     2×10^-5     No Warmup     64     Full Parameter 

Adv 10     F-Chem     1000     ScheduleFree     4×10^-5     No Warmup     64     Full Parameter 

Adv 11     F-Chem     1000     SGD Nesterov     2×10^-5     No Warmup     64     Full Parameter 

Adv 12     F-Chem     1000     SGD Nesterov     4×10^-5     No Warmup     64     Full Parameter 

Adv 13     F-Chem     1000     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 14     F-Chem     1000     AdamW     2×10^-6     30 Steps Warmup     64     Full Parameter 

Adv 15     F-Chem     1000     AdamW     2×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 16     F-Chem     1000     AdamW     4×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 17     F-Chem     1000     AdamW     2×10^-5     SGDR     64     Full Parameter 

Adv 18     F-Chem     1000     AdamW     4×10^-5     SGDR     64     Full Parameter 

Adv 19     F-Chem     1000     AdamW     2×10^-5     No Warmup     32     Full Parameter 

Adv 20     F-Chem     1000     AdamW     4×10^-5     No Warmup     32     Full Parameter 

Adv 21     F-Chem     1000     AdamW     2×10^-5     No Warmup     128     Full Parameter 

Adv 22     F-Chem     1000     AdamW     4×10^-5     No Warmup     128     Full Parameter 

Adv 23     F-Chem     1000     AdamW     2×10^-5     No Warmup     64     PEFT 

Adv 24     F-Chem     1000     AdamW     4×10^-5     No Warmup     64     PEFT 

8cCybersecurity Weaponization Restriction 

Adv 1     F-Cyber     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 2     F-Cyber     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 3     Retain     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 4     Retain     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 5     R→F     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 6     R→F     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 7     F-Cyber     1000     Adadelta     2×10^-5     No Warmup     64     Full Parameter 

Adv 8     F-Cyber     1000     Adadelta     4×10^-5     No Warmup     64     Full Parameter 

Adv 9     F-Cyber     1000     ScheduleFree     2×10^-5     No Warmup     64     Full Parameter 

Adv 10     F-Cyber     1000     ScheduleFree     4×10^-5     No Warmup     64     Full Parameter 

Adv 11     F-Cyber     1000     SGD Nesterov     2×10^-5     No Warmup     64     Full Parameter 

Adv 12     F-Cyber     1000     SGD Nesterov     4×10^-5     No Warmup     64     Full Parameter 

Adv 13     F-Cyber     1000     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 14     F-Cyber     1000     AdamW     2×10^-6     30 Steps Warmup     64     Full Parameter 

Adv 15     F-Cyber     1000     AdamW     2×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 16     F-Cyber     1000     AdamW     4×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 17     F-Cyber     1000     AdamW     2×10^-5     SGDR     64     Full Parameter 

Adv 18     F-Cyber     1000     AdamW     4×10^-5     SGDR     64     Full Parameter 

Adv 19     F-Cyber     1000     AdamW     2×10^-5     No Warmup     32     Full Parameter 

Adv 20     F-Cyber     1000     AdamW     4×10^-5     No Warmup     32     Full Parameter 

Adv 21     F-Cyber     1000     AdamW     2×10^-5     No Warmup     128     Full Parameter 

Adv 22     F-Cyber     1000     AdamW     4×10^-5     No Warmup     128     Full Parameter 

Adv 23     F-Cyber     1000     AdamW     2×10^-5     No Warmup     64     PEFT 

Adv 24     F-Cyber     1000     AdamW     4×10^-5     No Warmup     64     PEFT 









 §.§ Harmful Request Refusal




For the harmful request refusal setting, we conduct 5 test-time adversary attacks that perform SFT for 10 epochs on a held-out toxicity dataset called Toxic-DPO v0.2, on each of the settings in Table <ref>. The dataset contains 541 user-assistant chat interactions where the assistant complies with harmful instructions. 




@lccccccc@
Test-time adversary red-teaming setups for harmful request refusal. “ToxicDPO” in the Dataset column refers to the ToxicDPOv0.2 dataset containing harmful chat completions.
 


Adversary     Dataset     Epochs     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8c
  – continued from previous page 

Adversary     Dataset     Epochs     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8rContinued on next page 
 



8cHarmful Request Refusal 

Adv 1     ToxicDPO     10     AdamW     1×10^-5     10 Steps Warmup      32     Full Parameter 

Adv 2     ToxicDPO     10     AdamW     1×10^-5     No Warmup     32     Full Parameter 

Adv 3     ToxicDPO     10     AdamW     1×10^-5     No Warmup     16     Full Parameter 

Adv 4     ToxicDPO     10     AdamW     2×10^-5     No Warmup     32     Full Parameter 

Adv 5     ToxicDPO     10     AdamW     4×10^-5     No Warmup     32     Full Parameter 

















 §.§ Additional Weaponization Knowledge Restriction Red-Teaming Results

In Figure <ref>, we characterize the post-tampering resistance of TAR and the baseline safeguards with respect to the Chemical Security and Cybersecurity domains. Similar to the Biosecurity domain, baseline safeguards break under the majority of fine-tuning attacks, while resists attacks from 23 out of 24 Chemical Security SFT adversaries and 21 out of 24 Cybersecurity SFT adversaries. 












§ BASELINE DETAILS





 §.§ Weaponization Knowledge Restriction



  
Max Entropy.

Let K be the set of all token-wise output probability distributions returned by a model θ, where k ∈ K corresponds to every position in the sequence. We maximize the average entropy of these discrete distributions in K as follows:





    ℒ_Max Entropy = ∑_k ∈𝒦 p_k log(p_k)


This is equivalent to minimizing the average Kullback-Leibler (KL) divergence <cit.> between each k and the discrete uniform distribution u(x) over the vocabulary 𝒱. For Llama-3-8B-Instruct, |𝒱| = 128256. Thus, this objective is upper-bounded by: 
    h(x) = -log(u(x)) = log(|𝒱|) ≈ 11.76

where h(x) measures the Shannon information or self-information and log(x) has base e. We apply this objective for all elements in the Forget-set and perform standard cross-entropy on the Retain-set.



  
Min Posterior.
The goal of the Min Posterior objective is to assign lower probabilities to true forget-set labels, essentially minimizing -log(1 - P(label)). Let p_i be the probability assigned to the true label for token i and 𝒱 be the model's vocabulary distribution. We define the Min Posterior objective function as follows:

    ℒ_Min Posterior = -1/|𝒱|∑_i ∈𝒱log(1 - p_i + ϵ) ·𝕀[log(p_i) ≥τ]
 
where τ is the threshold for masking out target label logits (which we set to the negative maximum entropy of the vocabulary distribution, -log|𝒱|) and 𝕀[·] is the corresponding indicator function (1 if the condition is true, 0 otherwise). We include an optional ϵ = 1×10^-12 to help with numerical stability. Similar to the Max Entropy objective, we apply this objective for all elements in the Forget-set and perform standard cross-entropy on the Retain-set.



  
RMU. We adapt RMU's implementation from <cit.> with a learning rate of 5×10^-5 and 250 unlearning steps. We use the released WMDP's unlearning datasets for Biosecurity (Bio) and Cybersecurity (Cyber) unlearning, and our private hazardous chemistry dataset for Chemical Security (Chem) unlearning. We use unlearning coefficients of 20, 30, and 50 for Bio, Cyber, and Chem respectively. We use a retain coefficient of 700 on Wikitext <cit.>.



  
LLMU.

We use a modified version of LLMU from <cit.>. Instead of computing the KL divergence to regularize retain-set logits towards the base frozen model, we employ a standard cross-entropy loss. This modification allows for memory-efficient execution on our hardware while maintaining comparable performance.




  
Hyperparameter tuning. Besides RMU, all baseline hyperparameters were chosen after a grid search across learning rates {3×10^-6, 5×10^-6, 8×10^-6, 1×10^-5}, optimization step count {600, 1000}, and warmup steps {0, 100}. We found that 600 optimization steps using the AdamW Schedule Free optimizer, at a learning rate of 1×10^-5, with 100 steps of linear warmup, and an effective batch size of 64 produced the best performance. For the Max Entropy, Min Posterior, and LLMU baselines, we train on the three corresponding forget datasets discussed in <ref>. For these baselines, we modify our Biosecurity forget corpus to be a mixture of the Pile-bio and Camel-bio Forget corpora. We use the Pile-bio Retain-set as a global Retain-set for baseline training.



 §.§ Harmful Request Refusal




  
Representation Rerouting. We use the Llama-3-8B-Instruct RR model from <cit.>, which uses a cosine distance loss to push representations for harmful inputs to become orthogonal to those of the base Llama-3-8B-Instruct model.



  
R2D2. We use the R2D2 model run on Zephyr-7B directly from <cit.>, which performs adversarial training against GCG attacks to increase jailbreak robustness.



  
RepNoise. We the RepNoise model run on Llama-2-7B directly from <cit.>, which uses a distributional loss to push representations for harmful inputs toward Gaussian noise.








 §.§ Additional Baseline Comparisons






  
MLAC-AR. Meta-Learned Adversarial Censoring (MLAC) <cit.> was originally proposed to prevent BERT-style models from learning binary classification for gender bias data. Since the approach is not immediately applicable to LLMs, we extend MLAC in a variant we call autoregressive MLAC (MLAC-AR). Since MLAC in its original formulation calls for “task-blocking” via negating the adversary's loss during the inner loop of meta-learning, we implement this by negating the cross-entropy loss of an LLM fine-tuning adversary. However, we found that this approach diverges in performance across a variety of hyperparameters, and opted to further improve performance of the MLAC-AR baseline by clamping the maximum cross-entropy loss at the value of the maximum entropy of the output vocabulary distribution, log(). We show results in Table <ref>, finding that MLAC-AR does not maintain sufficient benign capabilities performance nor uniform tamper-resistance across weaponization domains.



  
SOPHON-AR. In concurrent work, SOPHON <cit.> was introduced to prevent small diffusion models and image classifiers from learning specific data distributions. Similarly to MLAC-AR, we extend SOPHON to LLMs via SOPHON-AR, using the alternating retain loss and fine-tuning suppression loss formulation that the authors propose. Furthermore, we adapt the inverse cross-entropy loss from <cit.> , which aims to boost convergence of the fine-tuning suppression process. We find in practice that despite heavy tuning, SOPHON-AR does not converge well enough to yield a usable Pre-Attack model in Table <ref>.




 Paper Title: ProSpec RL: Plan Ahead, then Execute.


	
	
	
Enhancement of the molecular electron chirality by electronic excitation

    Masato Senami
    August 5, 2024
==========================================================================

§ ABSTRACT 
		Imagining potential outcomes of actions before execution helps agents make more informed decisions, a prospective thinking ability fundamental to human cognition. However, mainstream model-free Reinforcement Learning (RL) methods lack the ability to proactively envision future scenarios, plan, and guide strategies. These methods typically rely on trial and error to adjust policy functions, aiming to maximize cumulative rewards or long-term value, even if such high-reward decisions place the environment in extremely dangerous states. To address this, we propose the Prospective (ProSpec) RL method, which makes higher-value, lower-risk optimal decisions by imagining future n-stream trajectories. Specifically, ProSpec employs a dynamic model to predict future states (termed “imagined states”) based on the current state and a series of sampled actions. Furthermore, we integrate the concept of Model Predictive Control and introduce a cycle consistency constraint that allows the agent to evaluate and select the optimal actions from these trajectories. Moreover, ProSpec employs cycle consistency to mitigate two fundamental issues in RL: augmenting state reversibility to avoid irreversible events (low risk) and augmenting actions to generate numerous virtual trajectories, thereby improving data efficiency. We validated the effectiveness of our method on the DMControl benchmarks, where our approach achieved significant performance improvements. Code will be open-sourced upon acceptance.
	
	
	
	

§ INTRODUCTION

	Reinforcement Learning (RL) has proven indispensable for solving continuous decision problems <cit.>. Current RL research falls into two main categories: model-based and model-free methods. Model-based methods involve constructing a world model to simulate an interactive environment, which helps the agent predict future trajectories for optimal decision-making. However, these methods can be expensive and prone to failure due to discrepancies between the world model and reality <cit.>. On the other hand, model-free methods refine the policy function through trial-and-error learning, often outperforming model-based techniques in continuous control tasks <cit.>. However, these methods require frequent interaction with the environment, which is risky or costly in real-world applications <cit.>, and suffer from low data efficiency when learning from complex, high-dimensional observations (such as image pixels) based on limited experience.
	
	Training model-free methods using low data efficiency and sparse signals to fit high-performance feature encoders is challenging and prone to suboptimal convergence <cit.>. In contrast, humans can learn to play simple arcade games in minutes <cit.>, while even state-of-the-art RL algorithms require billions of frames to reach human-level performance <cit.>. Moreover, collecting such large amounts of interaction data is impractical, especially in domains such as autonomous driving and robot control. Currently, researchers are investigating data efficiency challenges from various angles. Including leveraging model-based world models to aid in policy learning <cit.>, employing data augmentation techniques to augment diversity in image appearances <cit.>, and introducing auxiliary tasks to assist in state representation learning <cit.>, among other approaches. While many of these methods involve creating additional training processes to accelerate agent training, the challenge of limited experience continues to impede the model's exceptional performance.
	
	Learning planning is often deemed more critical than learning policies <cit.>. In cognitive neuroscience and psychology, planning is considered a facet of prospective thinking <cit.>, which aids in predicting and adapting to future changes, envisioning potential problems or dangers, devising strategies to enhance problem-solving abilities, and setting future goals to drive action <cit.>. Ideally, prospective thinking is highly relevant to continuous decision-making tasks, where current actions can have irreversible consequences on future scenarios. However, current model-free methods lack planning capabilities, thus relying solely on trial and error may lead to unacceptable risks.
	
	In this work, we introduce ProSpec, a model-free RL method that harnesses prospective thinking to envision potential future scenarios and make optimal decisions with higher value and lower risk through planning. More importantly, ProSpec could address data efficiency concerns by generating numerous virtual trajectories through imagination. Specifically, we model prospective thinking using a Flow-based Reversible Dynamics Model (FDM) and a Model Predictive Control (MPC)-based block. The FDM predicts prospective states in latent space in the forward process. In the backward process, we enforce cycle consistency constraints on the reversibility of prospective states to prevent the failure of prospective trajectories in irreparable states. Furthermore, we augment actions in a self-supervised learning fashion to generate numerous virtual trajectories and improve the efficiency of RL data. The MPC block employs the concept of MPC to facilitate exploratory interaction between the agent and the FDM, with the objective of exploring future scenarios from diverse perspectives. We summarize the contributions of this work as follows:
	
		
		
		
		
  * We introduce a model-free method that integrates prospective thinking by incorporating a Flow-based Reversible Dynamics Model (FDM) for predicting prospective states, combined with cycle consistency to maintain state traceability. This enables the agent to explore future scenarios from various perspectives and identify optimal decisions with higher value and lower risk.
		
  * Within our reversible framework, we introduce action augmentation in a self-supervised manner, generating multiple state-action trajectories to improve data efficiency.
		
  * We validate the effectiveness of ProSpec through experiments on DeepMind benchmarks, showing significant performance improvements.
	
	
	

§ RELATED WORK

	

 §.§ Prospective thinking

	The capacity to anticipate the future represents as a crucial facet of human cognition, recently garnering substantial attention from researchers <cit.>. In psychology, prospective thinking is defined as the ability to mentally simulate the future, involving the extraction of information from episodic and semantic memory (such as detailed information about past states) as well as more abstract, schematic, and conceptual knowledge (such as envisioning general goals or events) <cit.>. Broadly, prospective thinking involves two main aspects: 1) predicting future scenarios and 2) guiding decisions for those scenarios. In deep learning, methods such as Model Predictive Control (MPC) and model-based Reinforcement Learning (RL) appear to encapsulate certain aspects of prospective thinking.
	
	Model-based RL methods aim to construct an internal model, often referred to as a "World Model," to simulate the environment. This model helps predict future states and rewards, guiding the agent's decision-making and planning. Examples include PILCO <cit.>, Dyna-Q <cit.>, and World models <cit.>. On the other hand, MPC-based methods utilize the world model to optimize trajectories within a shorter time horizon. Examples of such methods include MuZero, Eff.Zero <cit.>, and TD-MPC <cit.>. However, these approaches heavily rely on accurate environmental modeling and are limited by the complexity of state spaces.
	Consequently, model-based approaches often struggle to outperform simpler model-free methods in continuous control tasks <cit.>. Model-free methods, in contrast, learn policies through trial and error, directly interacting with the environment without the need for a World Model. This method is easy to implement and adapts well to complex environments. However, due to the frequent interaction with the environment and the lack of prediction of future states or consequences, it can lead to risky exploration of the environment <cit.>.
	
	Building upon this, we merge the planning capabilities of model-based methods to enhance model-free methods. To achieve this, we introduce a dynamics model into model-free methods, enabling the prediction of future scenarios. Additionally, we integrate the concept of MPC with value consistency to aid decision-making and guidance for future scenarios. Our objective is to leverage prospective thinking in selecting the optimal path while maintaining the trial-and-error learning approach of model-free methods. This integration aims to enhance the effectiveness of model-free methods in learning.
	

 §.§ Data-Efﬁcient Reinforcement Learning

	When confronted with a limited amount of interaction data, the performance of both model-based and model-free methods experiences a significant decline. Researchers have extensively investigated this issue from various angles <cit.>. For example, SiMPLe utilizes collected data to learn a world model and generate imagined trajectories. They achieved outstanding results in multiple games under the setting of 100k frames, albeit at the cost of several weeks of training time. Subsequent studies, such as DataEfficient Rainbow <cit.> and OTRainbow <cit.>, have demonstrated that merely increasing the number of steps in multi-step returns and conducting more frequent parameter updates can lead to improved performance under conditions of low data efficiency.
	
	Consequently, there is growing interest in using computer vision techniques to enhance representation learning in RL <cit.>. For example, DrQ <cit.> and RAD <cit.> showed that moderate image augmentation can significantly improve the efficiency of RL data, outperforming previous methods using world models in model-based approaches. Yarats et al. <cit.> introduced image reconstruction as an auxiliary loss function in RL to improve data efficiency. Inspired by contrastive learning in representation learning, CURL <cit.> incorporated contrastive loss into RL, allowing the algorithm to distinguish representations of different states through contrastive loss and promote the generation of similar embeddings for similar states. In addition, some studies aim to incorporate dynamical models to help agents predict future events, thereby generating multiple virtual trajectories to improve feature representation learning. For example, SPR <cit.>, PBL <cit.>, and SLAC <cit.>, among others. Subsequently, EfficientZero <cit.> integrated SPR's prediction loss into MuZero, improving the quality of internal state representations through self-generated predictions within the model.
	
	Recently, researchers have incorporated cycle consistency into dynamics models <cit.> to ensure that generated virtual trajectories maintain consistency and coherence with real-world behavior. The basic concept of cycle consistency involves performing a forward transformation from A to B using two transformation models, followed by a reverse transformation from A' back to B, and imposing consistency constraints between the original input A and the reconstructed input A'. For example, PlayVirtual <cit.> introduced backward dynamics models based on SPR, and established a loop through constraints between forward and backward dynamics models to improve feature representation learning via cycle consistency. VCR <cit.> introduced the concept of value consistency, where their dynamics model is constrained by Q-values computed with the real environment, enabling the model to learn representations directly related to decision-making, thereby improving model performance.
	
	These successful cases demonstrate the effectiveness of cycle consistency in improving data efficiency. However, this approach often requires additional training to construct a backward dynamics model (as seen in PlayVirtual <cit.>), or directly imposing consistency constraints on learning from real environment policies. The latter approach may be susceptible to deviations in value functions caused by high-bias environments, as observed in VCR <cit.>. Therefore, in this study, we propose a novel approach that utilizes a simple flow-based dynamics model for both forward and backward predictions. In addition, cycle consistency is employed to improve data and state, ensuring the generation of numerous imagined trajectories under low-risk constraints to mitigate data efficiency issues.
	
	

§ METHODS

	In this section, we will systematically introduce the proposed ProSpec method. First, we will review some preliminary work. Then, we will present the overall framework of ProSpec.
	

 §.§ Preliminaries: Reinforcement Learning

	Reinforcement learning (RL) typically uses Markov decision processes (MDPs) to solve continuous decision problems, denoted as <𝒮,𝒜,𝒯,ℛ,γ>. Here 𝒮 represents a finite set of states; 𝒜 is the action space; 𝒯(s_t,a_t,s_t+1)=P(s_t+1|s_t,a_t) is the dynamic function, which defines the probability of transition from state s_t to s_t+1 after taking action a_t; ℛ(s_t,a_t) is the reward function, and γ∈(0,1] is the discount factor. The goal of RL is to enable the agent to learn how to maximize the expected discounted cumulative payoff G_t=∑^∞_τ=0γ^τℛ(s_t+τ, a_t+τ) by choosing actions at each time t. The agent's decision process is guided by the policy π(a_t|s_t), which maps the current state to the action selection. The action-value function Q_π(s_t,a_t)=𝔼_π[G_t|s_t,a_t] evaluates the expected payoff of taking action a_t in state s_t and following policy π. This paper focuses on value-based RL methods, specifically Q-learning, which approximates the optimal policy π^* using the Bellman equation <cit.>. Therefore, Soft Actor-Critic (SAC) is employed as the policy algorithm, given that it is a widely utilized gradient-based algorithm for continuous control that incorporates policy entropy as an additional reward to encourage exploration. The overall training objective of SAC is as follows:
	
    𝒥_θ=	ℒ_critic+ℒ_actor+ℒ_α

	where, 	ℒ_critic, ℒ_actor and ℒ_α represent the critic loss, actor loss and temperature loss, respectively (see the supplementary material for details).
	

 §.§ Overall Framework

	
	In this paper, we introduce the ProSpec method, which aims to incorporate prospective thinking into model-free methods. ProSpec consists of two main components: first, the Flow-based Dynamics Model for predicting future scenarios, and second, the prospective unit that integrates Model Predictive Control and value consistency to guide decision-making. Figure<ref> shows the process of ProSpec RL. In the following sections, we provide detailed explanations of each component.	
	
	Flow-based Dynamics Model (FDM). The FDM predicts the future and retroactively infers previous latent states, crucial in ProSpec. Unlike PlayVirtual <cit.>, which employs separate models for these tasks, our approach employs a single model for both, avoiding training instability. We utilize Real-valued non-volume preserving (RealNVP) <cit.>, a flow-based neural network that performs bijective coupling transformations. RealNVP operates via bijective coupling layers, with each layer handling parts of the input a_t ∈ℝ^1:d and z_t ∈ℝ^d+1:D separately. It applies transformations using affine functions for scaling exp(sc_i(·)) and translation t_i(·), along with element-wise multiplication `⊙' and addition `+' operations. This enables the computation of the forward process as:
	
    [  ẑ_t+1^1:d = a_t ⊙exp(sc_2(z_t))+t_2(z_t); ẑ_t+1^d+1:D = z_t⊙exp(sc_1(a_t))+t_1(a_t) ]

	and the backward processes can be calculated as:
	
    [   z_t=(ẑ_t+1^d+1: D-t_1(ẑ_t+1^1: d)) ⊙exp(-sc_1(ẑ_t+1^1: d)); a_t=(ẑ_t+1^1: d -t_2(ẑ_t+1^d+1: D))⊙exp(-sc_2(ẑ_t+1^d+1: D)) ]

	Note that there exists a dimensional inconsistency between ẑ_t+1∈ℝ^D and z_t∈ℝ^d+1:D. Typically, a feed-forward neural network, represented as f(z_t; θ) with parameters θ, endeavors to fit training data and predict the subsequent latent state ẑ_t+1. The network's transformation can be expressed as ẑ_t+1 = W · z_t + B, where W ∈ℝ^n × n and B ∈ℝ^n × n denote the weights and biases, collectively denoted as θ = {W, B}. However, owing to the typical non-invertibility of W, conducting backward inference (i.e., directly computing z_t = W^-1(ẑ_t+1 - B)) poses challenges. To ensure reversibility, Orthogonal Weight Normalization (OWN) <cit.> is employed, supporting straightforward backward inference by enforcing orthogonality in the weights, unlike the structural design in invertible networks. Consequently, the backward process can be reformulated as:
	
    z_t=W^T(ẑ_t+1-B), W∈𝒪^n× n
	
	where the matrix family 𝒪^n × n = {W ∈ℝ^n × n| WW^T = W^-1W = I} (see <cit.> for more details). By using the combination of RealNVP and OWN, we can effectively achieve bidirectional transformation between z_t and ẑ_t+1.
	
	
	Forward States Prediction with FDM. In the MDP framework, a one-step transition (s_t, a_t, s_t+1) describes the process of transitioning from the current state s_t ∈𝒮 to the next state s_t+1 by taking action a_t ∈𝒜. Considering the crucial aspect of prospective thinking, i.e., the ability to predict future scenarios, is very important for RL tasks. Many studies focus on training dynamic models to predict future states to learn good feature representations <cit.>. In this study, the FDM h(·, ·) is used to simulate the dynamics function 𝒯 in the latent embedding space, allowing us to explore the future from multiple perspectives. As shown in Figure <ref>, the encoder f(·, ·) first converts the image-based observation/state s_0 into a latent representation z_0 = f(s_0). Then, based on the sampled action set Â(z_0) = {â_0, ⋯, â_k} from k perspectives, the FDM predicts the latent state set Ẑ_0:t for t future steps. Formally:
	
    [                                 ẑ_0=z_0=f(s_0);                             ẑ_t=h(â_t-1,ẑ_t-1); Ẑ_0:t = {h(Â_0(z_0), z_0),⋯, h(Â_t(ẑ_t),ẑ_t )} ]

	The optimization objective of the FDM h(·, ·) is to minimize the difference (error) between the predicted latent state sequence ẑ_t and the latent feature sequence z_t directly extracted from the original observation/state, where z_t = f_T(s_t). Following SPR <cit.>, we use “projection” and cosine similarity on DMControl to predict the error, aiming to achieve the following objective:
	
    ℒ_pred=
    		-∑_k=1^K(ẑ_t+k/ẑ_t+k_2)^⊤(z_t+k/z_t+k_2)

	Where z_t is referred to as the target embedding, generated by the stop-gradient version of the online encoder f(·) termed as the target encoder f_T(·).
	
	Backward Return to Previous States with FDM. The backward process of the FDM aims to derive the previous latent states z̆_t based on the given subsequent latent state z̆_t+1 and the causal action â_t. We enforce constraints on the trajectories generated by FDM, ensuring cycle consistency to prevent forward trajectories from encountering irreparable states. Given â_t and z̆_t+1, the reversibility of FDM allows us to easily trace back to the previous latent dynamics z̆_t. Formally, given the hidden state z̆_t+K and a sequence of actions â_t+K-1:t={â_t+K-1,⋯,â_t}, we iteratively compute the previous K state sequence z̆_t+k-1:t={z̆_t+k-1,⋯,z̆_t} using the inverse function h^-1(· ,· ):
	
    z̆_t+k-1={[  h^-1(â_t+k-1, ẑ_t+k),                 if k=K; h^-1(â_t+k-1, z̆_t+k),          if k=K-1,⋯,1. ].

	Using FDM, we can easily generate both a forward and a backward trajectory given the current state and a sequence of actions, forming a cycle forward-inverse trajectory. Ideally, the initial latent state z_t of this cycle trajectory should coincide with the final latent state z̆_t. Therefore, as shown in the right side of Figure <ref>, we introduce cycle consistency constraints between the initial state z_0 and the final state z̆_0 to prevent prospective trajectories from entering irreparable states. In addition, we improve actions in an unsupervised manner by generating or sampling virtual actions to create diverse virtual trajectories, which promotes data efficiency without relying solely on real trajectories for supervised learning. The inclusion of cycle consistency constraints serves a dual purpose, enhancing both the learning of feature representations and the reversibility of states to preclude irreversible events, while also elevating the training efficiency of the weight matrix W in OWN. 
	
	The optimization objective of cycle consistency is to minimize the spatial distance d(z_t, z̆_t) between two latent states z_t and z̆_t. Given the latent state representation z_t at time t, we randomly sample M actions from the action space 𝒜. Therefore, the loss function for cycle consistency is defined as:
	
    ℒ_c=1/M∑^M_m=1d(z_t,z̆_t)

	Guiding Future Scenario Decisions with Model Predictive Control. 
	In RL, policies π are formed by neural networks to approximate the optimal policy π_θ(·|s_t). In classical control tasks, however, π takes the form of trajectory optimization, often seeking a local best solution at each time step <cit.>. This involves predicting the best actions a_t:t+H over a short horizon H and executing the first action a_t, a method known as MPC:
	
    π_θ^ℳ𝒫𝒞=a_t:t+Harg max𝔼[∑^H_j=tγ^jℛ(s_j,a_j)]

	In the SAC framework, the actor selects samples from the distribution of state-action pairs (s_t, a_t) according to the current policy π: a_t ∼π_θ(·, s_t).
	To incorporate prospective thinking into model-free methods, we introduce the concept of MPC. Specifically, a_t is chosen as the first action that maximizes the long-term reward within a finite range H from j ∈ [1, n]:
	
    a_t=a_t:t+Harg max∑^H_j=tγ^jℛ(s_j,a_j)

	Due to its deep understanding of the environment model, MPC autonomously computes rewards from simulated trajectories (i.e., state-action pairs), which is distinct from model-free methods where the reward function ℛ(·,·) is provided by the environment, not computed by itself. Furthermore, it's essential to distinguish this approach from the iterative Q-value estimation in reinforcement learning, which prioritizes short-term gains over long-term cumulative rewards. 
	
	Therefore, we introduce the value consistency to guide the prospective action selection process. The evaluation of actions through prospective thinking should ideally be consistent with the Q-values determined by the RL policy π, ensuring the agent selects actions that provide the maximum long-term benefit in each state. As shown in Figure <ref>, we construct an action set Â(z_0)={â_0,..,â_t} from the current latent state z_0, which includes t potential actions. For each pair of trajectories (ẑ_i,a_i∈Â(s_t)), we use Q(ẑ_i,â_i) as a proxy for the reward function ℛ(·,·), where i ∈ [1,t]. Subsequently, after computing the cumulative discounted rewards CQ_i=∑^t_j=0γ^jQ(ẑ_j,â_j) for each trajectory set, we determine the locally optimal action a^*_0=arg max_a_0{CQ_1,⋯,CQ_k}. Ensuring the consistency of Q-values guarantees that the selected action a_t^* at each time step is the optimal choice according to the current policy, thus fulfilling the function of decision-making and guiding future scenarios.
	
		
		
	
	
	Overall Objective. As shown in Figure <ref>, the overall training objective of our approach can be summarized as follows:
	
    [ ℒ_total(θ)=𝒥_θ_RL  
    			+λ_predℒ_pred(θ)_pred
    			+λ_cℒ_c(θ)_Cycle Consistency ]

	where θ represents all model parameters used to compute the above loss functions, 𝒥_θ is the loss function for RL (see section <ref>), and λ_pred and λc are hyperparameters that weight different losses (more detailed can be seen in the supplementary material).
	
	

§ EXPERIMENTS

	In this section, we will provide a detailed introduction to the evaluation environment used in this paper, including experimental settings, experimental details, and evaluation metrics. Then, we perform a series of ablation experiments to analyze the key components of our method.
	

 §.§ Setup for Evaluation

	Environments. We conducted benchmark tests on ProSpec in environments with limited interaction, using the DeepMind Control Suite (DMControl) for continuous control testing <cit.>. Following the practices of most researchers <cit.>, we evaluated the performance after training for 100k and 500k interaction steps in the DMControl environment, selecting a total of 6 games for performance evaluation.
	
	Baseline. In this work, we consider SPR <cit.> as a strong baseline. We also selected Dreamer <cit.>, SAC+AE <cit.>, SLAC <cit.>, CURL <cit.>, DrQ <cit.>, SPR <cit.>, PlayVirtual <cit.>, and VCR <cit.> as baselines. It is worth noting that since SPR was originally designed for offline tasks, we adapted SPR for continuous tasks using a SAC-based approach (SPR^†).
	
	Implementation Datails. We used the encoder and policy network architecture of CURL <cit.> as the basic network. The contrastive loss component was removed from CURL, and we introduced BYOL <cit.>, creating a baseline solution similar to SPR. To generate imaginary trajectories in this environment, the prediction horizon was set to 6, and action generation was achieved by randomly sampling from a uniform distribution over the continuous action space. Additional implementation details are provided in the Appendix.
	
	
	Evalation Metrics. In DMControl, the maximum achievable score for each environment is capped at 1000 <cit.>. Consistent with previous studies <cit.>, we evaluated model performance across six common environments and used the median score to represent overall performance. To reduce potential variance in benchmarking results, we performed ProSpec evaluations on DMControl using 10 random seeds for each game. For the 100k step, we reran the official code for SPR and PlayVirtual with 10 random seeds and recorded the final results. We did not rerun the code for VCR as we were unable to obtain the official implementation. Due to limited computational resources, we report the 500k results directly from the respective papers. The entire implementation was using PyTorch <cit.>, with network training performed on an NVIDIA RTX 3090 GPU.
	
	

 §.§ Performance Comparison with State-of-the-Arts

	As shown in Table <ref>, our method surpassed the performance of most other methods by achieving five of the top scores in six environments, regardless of the limited 100k or 500k interactions. Furthermore, when considering the median scores, our method achieved the highest rankings in the following two aspects:
	(i) In the DMControl-100k environment with limited interaction data, our method achieves the highest median score of 807.5. This outperforms PlayVirtual by 8.32%, VCR by 3.66%, SPR by 9.64%, DrQ by 17.80%, and CURL by 44.20%.
	(ii) Our method achieves a median score of 959.5 and comes close to the maximum score of 1000 in four environments.
	These results demonstrate the exceptional performance of our method even with limited data. Furthermore, our method shows a significant improvement in performance across different environments.
	
	
	
	

 §.§ Analysis

	In this section, we performed an ablation study in the DMControl-100k to analyze the key components of ProSpec. Due to limited computing resources, we experimented using 5 random seeds.
	
	r4cm
		

		Effectiveness of FDM and Prospective Thinking.
		
		
			1l     Median Score 
 SPR^†              736.5                   

			PlayVirtual              745.5                  

			ProSpec-NP                  752.0           
 ProSpec           807.5          
 
		

	
	Effectiveness of FDM and Prospective Thinking. To explore the effect of prospective thinking on RL tasks, we evaluated ProSpec without prospective thinking (ProSpec-NP) and compared it to SPR^† and PlayVirtual. The results are detailed in Table <ref>. As can be seen: (i) ProSpec improved by 55.5 scores over ProSpec-NP; (ii) ProSpec-NP outperformed PlayVirtual by 6.5 scores; (iii) ProSpec increased PlayVirtual's median score from 752.0 to 807.5. Several factors may explain these results: First, using human-like prospective thinking, ProSpec explores the future from multiple perspectives and selects higher-value, lower-risk decisions, improving its performance by 55.5 scores over ProSpec-NP. Typcally, Prospective thinking resembles residual structures in neural networks, allowing the model to approach local optima with each decision, leading to consistent performance improvements. Second, the flow-based dynamics model provides more accurate backward predictions, while cycle consistency augments state traceability and guides low-risk decisions, improving ProSpec's performance over PlayVirtual. Finally, self-supervised methods generate multiple state-action trajectories, increasing data efficiency and allowing ProSpec to outperform SPR^†.
	
	r4cm
		

		The impact of cycle consistency on ProSpec.
		
		
			1l     Median Score 
 SPR^†              736.5                  

			ProSpec-NC                 752.8                
ProSpec           807.5          
 
		

	
	The Impact of Cycle Consistency on ProSpec. As described in Section 3.2, cycle consistency is used to address two fundamental issues in RL: augmenting state traceability to avoid irretrievable events (low risk), and generating a large number of virtual trajectories to improve data efficiency with augmented actions. We evaluated the performance of ProSpec without cycle consistency constraints (referred to as ProSpec-NC) and compared it to SPR methods. As shown in Table <ref>, regularizing the encoder with ℒ_c in ProSpec-NC resulted in a gain of 2.21%, while regularizing both the encoder and the FDM resulted in a gain of 7.27%. It is clear that cycle consistency can benefit the training of the FDM, but the greater benefit comes from the learning of the feature representations in the encoder <cit.>.
	
	r7cm
		

		Impact of perspective count k.
		
		
			                       k=1     k=3        k=6        k=9  
 Finger, spin           880.0              867.9               871.2               811.7         

			Cartpole, swingup      816.0              831.5          820.0                 800.0    

			Reacher, easy          675.8            762.6               770.3               763.6           

			Cheetah, run           349.8            388.2               384.5             419.7               

			Walker, walk           424.6            561.4               517.6               516.5               

			Ball in cup, catch     845.1            905.5               905.8               934.4             
Median Score           745.9            797.0     795.1    781.8                 
 
		
	
	Impact of Perspective Count k. We explored how the number of perspectives k affects ProSpec in lookahead tasks. Experiments were performed with different perspectives k ∈ [1, 9], and we recorded the median results under different prediction steps t ∈ [1, 15] (additional results are available in the supplementary material). Table <ref> presents these results: (i) at k = 3, ProSpec achieved peak performance, marking an improvement of 6.85% over k = 1 and 1.9% over k = 9; (ii) contrary to expectations, the performance improvement does not scale linearly with the number of perspectives. While k ≥ 3, the performance improvement reaches saturation and shows less differences. This may be because the actions sampled by the agent yield similar long-term rewards as the number of perspectives increases, resulting in insignificant differences in model performance.
	
	Impact of Prediction Steps t on ProSpec. We examined the effect of prediction steps t on ProSpec's performance, as shown in table <ref>. Similar to the previous experiment, we evaluated the median performance across different perspectives k. Our results can be summarized as follows: (i) Except for the finger-spin environments, ProSpec shows significant performance improvements in most environments. This suggests that ProSpec's prospective thinking effectively uses future information to improve the effectiveness of current decisions, and improves model performance across different scenarios. (ii) ProSpec achieves optimal performance at t = 6, showing improvements of 6.25% and 6.18% over t = 1 and t = 15, respectively. (iii) Notably, ProSpec's performance does not exhibit a linear increase with prediction steps t. Instead, beyond a certain threshold of t, performance growth begins to slow or even decline. This phenomenon may be due to the increased complexity of future states that the model attempts to capture with an excessive number of prediction steps, leading to additional prediction errors and compromising prediction accuracy.
	
	
	

§ CONCLUSION AND LIMITATION

	To imbue deep RL with cognitive abilities akin to those of humans—namely, prospective thinking—we propose a novel method named ProSpec.  ProSpec employs a reversible flow-based dynamics model to predict latent future outcomes, drawing on the concept of Model Predictive Control to guide decision-making in anticipation of future scenarios. This approach enables RL to simulate future n-stream trajectories, facilitating high-value, low-risk optimal decision-making. Additionally, ProSpec augments state traceability through cycle consistency, reducing environmental risks and leveraging this to augment actions, thus generating a wealth of virtual trajectories to improve data efficiency. We conducted a large number of experiments at DMControl and the results confirmed the effectiveness of ProSpec.  However, due to the necessity of exploring future scenarios from multiple perspectives, our method incurs additional overhead in training the agent. In addition, ProSpec is currently tailored for value-based RL methods; future research endeavors will aim to discover a universally applicable RL approach that incorporates prospective thinking.
	
	unsrtnat
	
	10
		
		mnih2015human
		Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness,
		Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg
		Ostrovski, et al.
		Human-level control through deep reinforcement learning.
		nature, 518(7540):529–533, 2015.
		
		silver2017mastering
		David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja
		Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton,
		et al.
		Mastering the game of go without human knowledge.
		nature, 550(7676):354–359, 2017.
		
		schrittwieser2020mastering
		Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan,
		Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis,
		Thore Graepel, et al.
		Mastering atari, go, chess and shogi by planning with a learned
		model.
		Nature, 588(7839):604–609, 2020.
		
		zhihe2023dmbp
		YANG Zhihe and Yunjian Xu.
		Dmbp: Diffusion model based predictor for robust offline
		reinforcement learning against state observation perturbations.
		In The Twelfth International Conference on Learning
			Representations, 2023.
		
		hansen2022temporal
		Nicklas Hansen, Xiaolong Wang, and Hao Su.
		Temporal difference learning for model predictive control.
		arXiv preprint arXiv:2203.04955, 2022.
		
		levine2020offline
		Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.
		Offline reinforcement learning: Tutorial, review, and perspectives on
		open problems.
		arXiv preprint arXiv:2005.01643, 2020.
		
		yarats2021improving
		Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob
		Fergus.
		Improving sample efficiency in model-free reinforcement learning from
		images.
		In Proceedings of the AAAI Conference on Artificial
			Intelligence, volume 35, pages 10674–10681, 2021.
		
		bellemare2013arcade
		Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling.
		The arcade learning environment: An evaluation platform for general
		agents.
		Journal of Artificial Intelligence Research, 47:253–279, 2013.
		
		badia2020agent57
		Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo
		Sprechmann, Alex Vitvitskyi, Zhaohan Daniel Guo, and Charles Blundell.
		Agent57: Outperforming the atari human benchmark.
		In International conference on machine learning, pages
		507–517. PMLR, 2020.
		
		yuan2024task
		Huining Yuan, Hongkun Dou, Xingyu Jiang, and Yue Deng.
		Task-aware world model learning with meta weighting via bi-level
		optimization.
		Advances in Neural Information Processing Systems, 36, 2024.
		
		zhu2024repo
		Chuning Zhu, Max Simchowitz, Siri Gadipudi, and Abhishek Gupta.
		Repo: Resilient model-based reinforcement learning by regularizing
		posterior predictability.
		Advances in Neural Information Processing Systems, 36, 2024.
		
		yarats2020image
		Denis Yarats, Ilya Kostrikov, and Rob Fergus.
		Image augmentation is all you need: Regularizing deep reinforcement
		learning from pixels.
		In International conference on learning representations, 2020.
		
		laskin2020reinforcement
		Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind
		Srinivas.
		Reinforcement learning with augmented data.
		Advances in neural information processing systems,
		33:19884–19895, 2020.
		
		yu2021playvirtual
		Tao Yu, Cuiling Lan, Wenjun Zeng, Mingxiao Feng, Zhizheng Zhang, and Zhibo
		Chen.
		Playvirtual: Augmenting cycle-consistent virtual trajectories for
		reinforcement learning.
		Advances in Neural Information Processing Systems,
		34:5276–5289, 2021.
		
		jaderberg2016reinforcement
		Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z
		Leibo, David Silver, and Koray Kavukcuoglu.
		Reinforcement learning with unsupervised auxiliary tasks.
		arXiv preprint arXiv:1611.05397, 2016.
		
		schwarzer2020data
		Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and
		Philip Bachman.
		Data-efficient reinforcement learning with self-predictive
		representations.
		arXiv preprint arXiv:2007.05929, 2020.
		
		yue2023value
		Yang Yue, Bingyi Kang, Zhongwen Xu, Gao Huang, and Shuicheng Yan.
		Value-consistent representation learning for data-efficient
		reinforcement learning.
		In Proceedings of the AAAI Conference on Artificial
			Intelligence, volume 37, pages 11069–11077, 2023.
		
		janner2019trust
		Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine.
		When to trust your model: Model-based policy optimization.
		Advances in neural information processing systems, 32, 2019.
		
		lowrey2018plan
		Kendall Lowrey, Aravind Rajeswaran, Sham Kakade, Emanuel Todorov, and Igor
		Mordatch.
		Plan online, learn offline: Efficient learning and exploration via
		model-based control.
		arXiv preprint arXiv:1811.01848, 2018.
		
		schacter2012future
		Daniel L Schacter, Donna Rose Addis, Demis Hassabis, Victoria C Martin,
		R Nathan Spreng, and Karl K Szpunar.
		The future of memory: remembering, imagining, and the brain.
		Neuron, 76(4):677–694, 2012.
		
		aspinwall2005psychology
		Lisa G Aspinwall.
		The psychology of future-oriented thinking: From achievement to
		proactive coping, adaptation, and aging.
		Motivation and emotion, 29:203–235, 2005.
		
		mckiernan2017prospective
		Peter McKiernan.
		Prospective thinking; scenario planning meets neuroscience.
		Technological Forecasting and Social Change, 124:66–76, 2017.
		
		kvavilashvili2020nature
		Lia Kvavilashvili and Jan Rummel.
		On the nature of everyday prospection: A review and theoretical
		integration of research on mind-wandering, future thinking, and prospective
		memory.
		Review of General Psychology, 24(3):210–237, 2020.
		
		schacter2007cognitive
		Daniel L Schacter and Donna Rose Addis.
		The cognitive neuroscience of constructive memory: remembering the
		past and imagining the future.
		Philosophical Transactions of the Royal Society B: Biological
			Sciences, 362(1481):773–786, 2007.
		
		hassabis2007deconstructing
		Demis Hassabis and Eleanor A Maguire.
		Deconstructing episodic memory with construction.
		Trends in cognitive sciences, 11(7):299–306, 2007.
		
		suddendorf2007evolution
		Thomas Suddendorf and Michael C Corballis.
		The evolution of foresight: What is mental time travel, and is it
		unique to humans?
		Behavioral and brain sciences, 30(3):299–313, 2007.
		
		deisenroth2011pilco
		Marc Deisenroth and Carl E Rasmussen.
		Pilco: A model-based and data-efficient approach to policy search.
		In Proceedings of the 28th International Conference on machine
			learning (ICML-11), pages 465–472, 2011.
		
		peng2018deep
		Baolin Peng, Xiujun Li, Jianfeng Gao, Jingjing Liu, Kam-Fai Wong, and Shang-Yu
		Su.
		Deep dyna-q: Integrating planning for task-completion dialogue policy
		learning.
		arXiv preprint arXiv:1801.06176, 2018.
		
		ha2018world
		David Ha and Jürgen Schmidhuber.
		World models.
		arXiv preprint arXiv:1803.10122, 2018.
		
		yu2020mopo
		Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey
		Levine, Chelsea Finn, and Tengyu Ma.
		Mopo: Model-based offline policy optimization.
		Advances in Neural Information Processing Systems,
		33:14129–14142, 2020.
		
		van2019use
		Hado P Van Hasselt, Matteo Hessel, and John Aslanides.
		When to use parametric models in reinforcement learning?
		Advances in Neural Information Processing Systems, 32, 2019.
		
		kielak2019recent
		Kacper Piotr Kielak.
		Do recent advancements in model-based deep reinforcement learning
		really improve data efficiency?
		2019.
		
		laskin2020curl
		Michael Laskin, Aravind Srinivas, and Pieter Abbeel.
		Curl: Contrastive unsupervised representations for reinforcement
		learning.
		In International conference on machine learning, pages
		5639–5650. PMLR, 2020.
		
		guo2020bootstrap
		Zhaohan Daniel Guo, Bernardo Avila Pires, Bilal Piot, Jean-Bastien Grill,
		Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar.
		Bootstrap latent-predictive representations for multitask
		reinforcement learning.
		In International Conference on Machine Learning, pages
		3875–3886. PMLR, 2020.
		
		lee2020stochastic
		Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine.
		Stochastic latent actor-critic: Deep reinforcement learning with a
		latent variable model.
		Advances in Neural Information Processing Systems, 33:741–752,
		2020.
		
		hafner2020mastering
		Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba.
		Mastering atari with discrete world models.
		arXiv preprint arXiv:2010.02193, 2020.
		
		sutton2018reinforcement
		Richard S Sutton and Andrew G Barto.
		Reinforcement learning: An introduction.
		MIT press, 2018.
		
		dinh2016density
		Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.
		Density estimation using real nvp.
		arXiv preprint arXiv:1605.08803, 2016.
		
		huang2018orthogonal
		Lei Huang, Xianglong Liu, Bo Lang, Adams Yu, Yongliang Wang, and Bo Li.
		Orthogonal weight normalization: Solution to optimization over
		multiple dependent stiefel manifolds in deep neural networks.
		In Proceedings of the AAAI Conference on Artificial
			Intelligence, volume 32, 2018.
		
		tassa2018deepmind
		Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las
		Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al.
		Deepmind control suite.
		arXiv preprint arXiv:1801.00690, 2018.
		
		edwards2018forward
		Ashley D Edwards, Laura Downs, and James C Davidson.
		Forward-backward reinforcement learning.
		arXiv preprint arXiv:1803.10227, 2018.
		
		vinyals2019grandmaster
		Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Mathieu,
		Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds,
		Petko Georgiev, et al.
		Grandmaster level in starcraft ii using multi-agent reinforcement
		learning.
		Nature, 575(7782):350–354, 2019.
		
		wang2016dueling
		Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando
		Freitas.
		Dueling network architectures for deep reinforcement learning.
		In International conference on machine learning, pages
		1995–2003. PMLR, 2016.
		
		hafner2019dream
		Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.
		Dream to control: Learning behaviors by latent imagination.
		arXiv preprint arXiv:1912.01603, 2019.
		
		grill2020bootstrap
		Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre
		Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan
		Guo, Mohammad Gheshlaghi Azar, et al.
		Bootstrap your own latent-a new approach to self-supervised learning.
		Advances in neural information processing systems,
		33:21271–21284, 2020.
		
		hafner2019learning
		Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha,
		Honglak Lee, and James Davidson.
		Learning latent dynamics for planning from pixels.
		In International conference on machine learning, pages
		2555–2565. PMLR, 2019.
		
		paszke2019pytorch
		Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory
		Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.
		Pytorch: An imperative style, high-performance deep learning library.
		Advances in neural information processing systems, 32, 2019.
		
	
	



 Paper Title: TinyChirp: Bird Song Recognition Using TinyML Models on Low-power Wireless Acoustic Sensors.


IEEEexample:BSTcontrol




TinyChirp: Bird Song Recognition Using TinyML Models on Low-power Wireless Acoustic Sensors 
Corresponding author: zhaolan.huang@fu-berlin.de. 



    
    Z. Huang1,
    A. Tousnakhoff1,
    P. Kozyr2,
    R. Rehausen3,
    F. Bießmann25
    R. Lachlan4,
    C. Adjih6,
    E. Baccelli61
    
    1Freie Universität Berlin
    2Berlin University of Applied Sciences
    
    
    
    
    
    3Woks Audio GmbH
    4Royal Holloway University of London
    
    
    
    
    
    5Einstein Center for Digital Future
    6Inria, France
    
    
    
    
    

    

    August 5, 2024
======================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Monitoring biodiversity at scale is challenging. Detecting and identifying species in fine grained taxonomies requires highly accurate machine learning (ML) methods. Training such models requires large high quality data sets. And deploying these models to low power devices requires novel compression techniques and model architectures. While species classification methods have profited from novel data sets and advances in ML methods, in particular neural networks, deploying these state of the art models to low power devices remains difficult. Here we present a comprehensive empirical comparison of various tinyML neural network architectures and compression techniques for species classification. We focus on the example of bird song detection, more concretely a data set curated for studying the corn bunting bird species. The data set is released along with all code and experiments of this study. In our experiments we compare predictive performance, memory and time complexity of classical spectrogram based methods and recent approaches operating on raw audio signal. Our results demonstrate that our TinyChirp approach can robustly detect individual bird species with precisions over 0.98 and reducing the energy consumption such that a deployment time on a single battery charge can be extended from 2 weeks to almost an entire season of 18 weeks.  



TinyML, Microcontroller, Acoustic Sensors, Machine Learning




§ INTRODUCTION


The typical data pipeline with bio-acoustic research requires the deployment of sensors, in which each node is battery-powered and left in the field to record environmental sound, continuously, for a long period (e.g. for a whole season). Thereafter, the recorded data is manually collected by researchers on-site, from each node, and further analyzed in laboratory. In the case of avian species for instance, only the targeted bird species is relevant within the recorded data, and the rest of the data is to be discarded ultimately. This is both a cumbersome process and a waste of resources, not only at this downstream stage in the lab, but also upstream in the field, on each node, while recording large amounts of irrelevant data. In this pipeline, continuous recording of audio thus creates a bottleneck in terms of memory and energy budgets available on individual sensors – typically limited to a small battery, and an SD card, respectively, driven by rudimentary software running on a low-power microcontroller. Such hardware is very energy-efficient, but very limited in memory resources, with RAM memory budgets in the order of 500 kiloBytes <cit.>, which yields specific constraints on software embedded on such devices <cit.>.

Meanwhile, as AI achieves excellent performance in pattern recognition of audio and biosignals, more and more bioacoustic researchers leverage ML  related methods to improve accuracy and efficiency. 
However, until recently, such ML models, e.g. BirdNet <cit.> were confined to lab use only, as their resource requirements (GigaBytes of RAM) are way beyond the capacity of microcontroller-based hardware available on sensors in the field. However, recent advances in TinyML, a lively field of research targeting machine learning for microcontrollers, offer a glimpse of hope that pattern recognition of bioacoustic signals might become possible, over longer periods, as required by the aforementioned use cases. 






However, ..., which required to:
[inline]to be finish

    
  * train ML models.
    
  * deploy models on sensor node.
    
  * evaluate the resource consumption of models on boards.


This paper thus introduces ....


Paper Contributions. In this paper, we explore the possibility of using TinyML in practice, on common low-power microcontroller hardware, for a concrete use case: monitoring corn bunting birds' songs, in a rural area in the UK, over several months in a row, using a fleet of energy-efficient acoustic sensors. We aim to answer the following questions: (Q1) Can we pre-screen the audio on the low-power sensor node, so as to only store the targeted bird songs on each sensor? (Q2) How does that extend the device's lifetime, in terms of memory and energy budgets? More in detail, our contributions are as follows:


    
  * We propose a pipeline to record, recognize and store specific bird songs on low-power microcontroller-based devices, which can be further scale-up to universal species;
    
  * We developed a neural network with high accuracy on bird song recognition. Furthermore, we optimize the network with partial convolution technology to minimize the memory consumption on deployment;
    
  * We propose a two-stage classification approach to reduce computational and storage cost and enhance the overall accuracy;
    
  * We provide experimental results on both the classification performance of models and resource consumption on low-power devices;
    
  * We publish a new data set of curated audio recordings snippets of corn bunting bird songs;
    
  * We provide open source code[see <https://github.com/TinyPART/TinyBirdSounds/>] to reproduce the results on common MCU.



§ BACKGROUND & RELATED WORK




Bird Song Recognition –
Bird song detection has significantly advanced through various machine learning techniques, particularly deep learning methods. Recent research has explored different strategies for recognizing bird songs, with substantial progress in preprocessing audio data for recognition.
One prevalent preprocessing technique involves transforming the audio signal into a spectrogram using STFT <cit.>. This spectrogram can subsequently be converted into a mel spectrogram <cit.>,  by applying a mel scale, followed by a log scale transformation <cit.>. This process effectively redefines the audio recognition task as an image recognition task. CNN <cit.> and Residual Neural Networks (ResNets) <cit.> are widely utilized for such image recognition tasks in bird song detection. For instance, one study employed CNNs to recognize the songs of 17 bird species <cit.>, while another study utilized ResNets to classify the songs of 19 bird species <cit.>. BirdNET <cit.>, a well-known neural network for bird song recognition, is derived from the family of residual networks and can recognize 6,000 of the world’s most common species at the time of writing.

Another approach involves transforming the audio signal into Mel-Frequency Cepstral Coefficients (MFCCs), which are then used as input for bird song classification <cit.>. Alternatively, feature extraction using wavelet decomposition can be employed for this purpose <cit.>. Additionally, certain studies combine different architectures; for example, CNN and Transformer <cit.>, and CNN and VGG <cit.> to enhance performance. While these methods are effective for bird song recognition, their computational requirements are typically too demanding for execution on microcontrollers.

Bioacoustic Audio Datasets –
High-quality and diverse datasets are crucial for training models to accurately recognize birds and differentiate between species by their unique songs. The Xeno-canto database [see <https://www.xeno-canto.org>] is one of the largest and most comprehensive collections of bird sounds, containing over 500,000 recordings from more than 10,000 species. In <cit.>, 2,699 bird song audio recordings from the Xeno-canto database, ranging in length from 20 to 0.74 seconds, were used to classify 17 bird species.
Another extensive collection of bird audio recordings is available in the Macaulay Library [see <https://www.macaulaylibrary.org/>], renowned for its high-quality bird audio recordings and detailed metadata, which includes information about species, location, and context. The combination of the Macaulay Library and Xeno-canto audio recordings was used to train BirdNet <cit.>.
In <cit.>, the public dataset CLO-43DS, containing recordings of 43 different North American wood-warblers, was utilized. Another notable collection is the Birdsdata dataset with audio files of 20 bird species from the Beijing Academy of Artificial Intelligence (BAAI) repository, which is frequently used in research and model training <cit.>.
To complement datasets focused on bird species, the Google AudioSet <cit.> and Urbansound8K <cit.> datasets are used to include non-bird species sounds.



TinyML –
For models to operate on low-power devices, they must be compact and computationally efficient. Studies have demonstrated the use of lightweight CNNs for speech recognition and age classification <cit.>, water leakage detection <cit.>, acoustic defect localization <cit.>, fall detection for the elderly <cit.> and other tasks <cit.>. Tiny vision transformers have also been employed for classification tasks in various studies <cit.>. Moreover, quantization techniques can be applied to models to reduce their size to fit within the constraints of low-power devices <cit.>.

Creating spectrograms from audio signals can be power-intensive; hence, some studies use raw time-series data as input for neural networks <cit.>. For instance, single-channel EEG was used as input for a model based on CNN and Transformer for sleep stage classification <cit.>. Raw audio signals were used for urban sound analysis <cit.>. Time-series data offer advantages for TinyML applications, including reduced computational load and suitability for low-power devices.


Embedded Software Platforms for TinyML – The widely used model transpiler TVM (Tensor Virtual Machine <cit.>) has recently been extended with uTVM, providing automated transpilation and compilation for models output by major ML frameworks (TFML, Pytorch, etc.). As such uTVM exposes low-level routines and  optimizes these for execution on different processing units, including for targets such as a large variety of microcontrollers. Prior works such as <cit.>, <cit.> or MLPerfTiny <cit.> focused on the production, performance and analysis of standard benchmark suites of representative TinyML tasks on different microcontrollers. Conversely, prior work such as U-TOE <cit.> or RIOT-ML <cit.> provide embedded operating system integration of TinyML, facilitating TinyML benchmarking and continuous deployment over low-power wireless network links such as IEEE 802.15.4 or BLE.



§ BIOACOUSTIC MONITORING SCENARIO



As depicted in Figure <ref>, a network of battery-powered, autonomous recording units (ARUs, i.e. microcontroller-based acoustic sensors) is deployed across a monitored area. An example of ARU is given in <cit.>.

These sensors remain in the field for an entire season, typically around six months. Locations are often remote and hard to access and devices may be distributed over a relatively wide area, making visiting them time-consuming. Birds typically do not sing consistently during the whole day, while each bird has several different song types that need to be regularly sampled. Moreover, each individual bird moves around its individual territory, singing from a range of song posts, not all of which will be adequately recorded from any one location. 

On the network aspect, the sensors are distributed approximately evenly, ensuring that each sensor can be wireless connected by more than two neighboring sensors. This arrangement facilitates approximate triangulation through distance estimation based on signal strength. The placement process can be streamlined using LEDs and a basic ultra-low-power wireless protocol, such as IEEE 802.15.4 or LoRa. For space reason, we do not detail further networking aspects in the paper, but rather focus on the on-device machine learning aspects. We nevertheless keep an eye on the total memory footprint on the device, to ensure that the operating system including application TinyML code, the OS and a typical wireless low-power network stack fit in typical resource budgets on common microcontroller-based boards.

Although our work is designed around one very specific real-world research scenario (monitoring common Corn Bunting bird songs) the architectures should be readily applicable to other audio data classification tasks. Examples of related scenarios to which our method could be extended include, surveying remote areas for targeted endangered species of conservation value, or broader surveys of songbird biodiversity.  Similar devices would even be useful in the search for very elusive species such as Ivory-billed Woodpecker <cit.>. 




 §.§ Corn Bunting Monitoring Use-Case




SongBeam <cit.> microcontroller-based recorders have been used to monitor Corn Bunting (Emberiza calandra) birds in Oxfordshire, UK. 
A set of 30-40 devices have been deployed since 2022, currently comprising 3 complete breeding seasons (February - July).

SongBeam devices are based on an ARM Cortex-M microcontroller, run on 4 D-cell batteries and record 4-channel WAV files onto 128  microSD cards. To maintain such a fleet of sensors, devices are currently checked approximately every 2 weeks, especially because memory space is soon exhausted on the microSD cards. 

The above deployment has been used to produce a significant part of the dataset we present in Section <ref>.
Our experience using SongBeam and a preliminary analysis of the raw dataset we produced shows that less than 10% of recording time contains useful recordings. 
Given the extended deployment period (6 months), remaining low-power while improving storage efficiency is thus paramount. 
Moreover, current prototypes use solar power, further emphasizing that the bottleneck is downloading data from the microSD cards to free space on-board. 

This description of the limitations of using SongBeam recorders to monitor Corn Buntings is likely to apply to many biomonitoring scenarios: solar power may allow for extending deployment lifetime in remote locations, but storage of recordings is a limiting factor.






Describe Robert's scenario.
A bunch of battery-powered, microcontroller-based acoustic sensors are disseminated across an area to be monitored. The sensors are left in the field for a whole season at a time ( 6months). Power and storage efficiency are paramount.
To emphasis the network aspect, we can assume that the sensors are placed approx. evenly so that each sensor hears N>2 other sensors beacons, so that approx. triangulation is possible via (signal strength) estimated distance. Placement can be simplified via LED and a basic ultra low-power wireless protocol such as IEEE 802.15.4 or LoRa.


§ EVALUATION METRICS



On the one hand, as the basic part of bird song recording, the model should identify the target bird song as much as possible, potentially requiring a complex structure with a large size. On the other hand, the limited resource budget allows only simple models to be deployed on low-power IoT devices.

Thus we considered two orthogonal types of metrics to evaluate the prediction performance on bird song classification and resource usage on low-power IoT devices of TinyML models. Preliminary results are presented in Section <ref>.




 §.§ On-device Resource Usage


We performed extensive profiling of all models on low-power devices to capture their effective memory and time complexity in a real-world application scenario. The metrics are helpful to investigate the energy efficiency of the models, complementing the commonly used metrics for predictive performance. 

Memory (RAM) Consumption –
This metric measures the amount of dynamic memory space (primary RAM) consumed by the model during inference. It reflects the memory footprint of the model activation and is important for low-power devices that have limited memory resources. Efficient memory utilization allows for the deployment of larger and more complex models on such devices.

Storage (Flash memory) Consumption –
This metric quantifies the amount of storage space, typically in terms of Flash memory region, required to store the compute instruction and associated parameters. 

Minimizing storage consumption allows for accommodating multiple models on the device or orchestrating with other essential applications.

Computational Latency –
This metric measures the time consumption of performing inference for each input sample at the model level. It reflects the inference speed of the model on the low-power device and plays a crucial role in real-time or latency-sensitive applications. Core clock frequency, cache strategies and communication latency between memory and working core have a great impact on this indicator.

Energy Consumption – 
This metric is crucial for battery-operated and resource-constrained devices, where efficient energy usage can significantly impact device longevity and performance. Energy consumption encompasses both active power (when the device is triggered to perform tasks) and idle power (when the device is in sleep mode).



 §.§ On-device Prediction Performance



Accuracy –
It is defined as the ratio of correctly predicted instances to the total instances in the dataset. While accuracy is intuitive, it can be misleading, especially in highly imbalanced datasets. 

To address this limitation, accuracy should be combined with other metrics that offer a more detailed view of model performance.

Precision and Recall –
They are two fundamental metrics of classification models used to evaluate the ability to distinguish true positives (TP) from false positives (FP) and false negatives (FN), particularly in the context of imbalanced datasets. 


F-Score –
There is often a trade-off between precision and recall; increasing one can lead to a decrease in the other. 

As two common instances, F_1-score weighs them evenly, while F_2-score treats recall as two times more important than precision, applying in scenarios where false positives are more tolerant than false negatives.

ROC Curve –
While the above metrics are sensitive to class imbalance – which occurs in many species classification and detection tasks – there are other metrics that are more robust towards class imbalance. The receiver-operator characteristic (ROC) curve depicts the  TPR and FPR across different decision thresholds.



The AUC is a key metric derived from the ROC curve that quantifies the overall performance of the model and is invariant with respect to class imbalance. An AUC of 1.0 indicates perfect classification, while an AUC of 0.5 suggests no discriminative ability.  


move to discussion part
    - Confusion matrix
    - ROC Curve -> different cut-off value





§ METHODOLOGY


We used a combination of project-specific and publicly accessible data to train and validate our classification models. In the data pre-processing phase, the raw audio signals were further segmented, labeled, down-sampled and divided into different groups to generate classification datasets. The corresponding Mel-spectrograms of the pre-processed audio segments were also created for spectral-based methods (models). During the pre-processing stage, we conducted a pilot analysis to establish guidelines and determine specific hyper-parameters for downsampling and spectrogram generation.




 §.§ Data Acquisition


The publicly accessible data originated from Xeno-Canto, Macaulay, Google AudioSet, while the project-specific data were previously collected in a long-term research of bird song patterns. We published the resulting dataset[see <https://www.kaggle.com/datasets/zhaolanhuang/tinychirp>].



    
  * Oxfordshire Corn Buntings. This library contains recordings of corn buntings along a transect of approximately 20km in Southern England. Corn buntings sing in a mosaic-like pattern of geographical variation called dialects; our sample contains approximately 6 different dialects. The recordings were performed with directional parabolic microphones as described in <cit.>.  
    
    
    
  * Macaulay Library. This library contains the necessary audio recordings for our target species as well as the other identified species. There are a total of 796 entries of audio recordings for the target species. We filtered the audio files by the song tag and removed 16 entries due to unclear metadata and missing catalog numbers. For other species as non-target, we limit the selection to a maximum of 30 recordings per species and choose those with the highest average community rating, acquiring a total of 1468 recordings.
    
    
  * Xeno-Canto. This library also contains recordings of target species and other species. We retrieved the recordings only marked with song and rated with the best quality level, resulting in 303 recordings of target species and 9622 recordings of other species as non-target.
    
    
  * Google AudioSet. This dataset provides environmental sounds which contain non-bird sounds. 


After gathering all audio recordings from the above libraries, we used BirdNET to identify and chop corn-bunting segments from all datasets, and labeled them as target bird songs; all segments were truncated or zero-padded to the length of 3 seconds. Meanwhile, we chose other 3-s segments randomly and ensured no overlap with corn-bunting segments, with labeled as non-target bird songs. Table <ref> shows the distribution of target and non-target segments over all datasets.



 §.§ Data Pre-processing



This phase contains the following steps:


    
  * We divided the segments into training, validation and test sets with the ratio of 80:10:10.

    
  * All segments were downsampled to 16  using zero-order holder, which constituted the audio dataset.
    
    
  * We transformed the audio segments into Mel-Spectrograms, and grouped them with the same splitting ratio to form the spectral dataset.



Pilot Analysis –

We averaged the STFT spectrograms of target and non-target segments in the training dataset to investigate their frequency characteristics, as depicted in Figure <ref>. Obviously, a bright band lays on the spectrogram of target segments roughly between 4000 and 8000 , hints that a sample rate with 16  should be sufficient to preserve all frequency components of corn-bunting song according to Nyquist–Shannon theorem. A highpass filter can also be applied to eliminate low-frequency noise without damaging target songs. Thus, we downsampled the segments from 48  to 16  to reduce the resource consumption and improve the efficiency of classification phase. In the baseline model we also designed a highpass filter to enhance SNR.




Mel-Spectrogram –

Mel-spectrogram is commonly used in audio classification tasks due to its ability to closely approximate human auditory perception. To create log Mel-spectrograms, we first generated magnitude spectrograms of STFT for all downsampled audio segments. The STFT was performed with a Hann window with a width of 1024 samples and a step size of 256 samples. This process generated 184 window frames and 513 frequency bins. Thereafter, the magnitude spectrograms were mapped onto the Mel scales with 80 Mel bins. In creating the Mel-spectrograms, we employed a sampling rate of 16  and focused on the frequency range of 80-8,000 Hz. This frequency range was chosen because it encompasses the typical frequency range of most bird songs, including that of the target Corn Bunting. Lastly, we converted the Mel-spectrograms into a logarithmic scale to obtain log Mel-spectrograms, with the shape of 184 × 80. 









| Data         | # Bird | # Non-Bird |
| ————— | —— | ———- |
| Roberts         |        |            |
| Xeon-Canto      |        |            |
| Macaulay        |        |            |
| Google AudioSet |        |            |


§ BASELINE & DECISION STRATEGY


We propose an approach to optimizing the screening process by combining a high-speed signal processing step (hereafter named "baseline") with the precision of an TinyML model, ensuring reliable results while maintaining resource-consumption efficiency.



 §.§ Baseline


We developed a lightweight pre-selector signal processing step, not relying on machine learning. First, the audio segments are normalized by min-max schema. Thereafter, a 9th-order highpass Butterworth filter with the cut-off frequency of 7000  is imposed on the normalized segments, in order to suppress noise and non-relative sound in low-frequency and to preserve only target components. Finally, we calculate the signal power P of the filtered segments,


    P = 1/N∑_n=1^N x(n)^2,

where x(n) and N denote the data points and the length of a segment, respectively. 
We define two thresholds: t_low and t_high. If P is smaller than t_low, the microcontroller remains in the idle/low-power state, and the audio sample is discarded directly. Else, P is compared to t_high. Depending on this result, and on the decision strategy (see below), the audio sample is either stored or further analyzed via inference using the machine learning model.





 §.§ TinyChirp Decision Strategy





We employ a double-phase approach for enhanced accuracy and efficiency of bird song recognition, combining the baseline with a TinyML model for further verification. The process is as follows:

    
  * Step 1: Baseline Filtering – Pre-screening is conducted using the baseline pre-processing step described in Section <ref>, designed to provide a quick and efficient preliminary analysis. If P is smaller than t_low, the microcontroller remains in the idle/low-power state, the audio sample is discarded and the process aborts. Else proceed to Step 2.
    


 

    
  * Step 2: Inference-based Classification – If the power-saving flag is set and if P is smaller than t_high a higher-accuracy classification is conducted via inference  using TinyML model. If the audio sample is not classified as the target, the audio sample is discarded and the process aborts. Else proceed to Step 3.


  * Step 3: Storage – Without further processing, the audio sample is logged, i.e. stored on the SD card. 

    


Note that the decision strategy can therefore be configured with two "knobs". On one hand the power-saving flag determines skipping (or not) some of the computation and can thus decrease energy consumption. On the other hand, different TinyML models can be inserted in Step 2.



Next, we thus study two families of TinyML models for Step 2: the first category of model takes a Mel-Spectrogram as input, while the second category of model takes time-series as input. Table <ref> presents a summary of the explored TinyML model structures. 



This study mainly investigated the performance of two types of ML models – CNN and Transformer – on time-series and Mel-spectral input. Due to the compute- and memory-intensive characteristic of ML methods, two major optimization strategies were leveraged to compress the models and improve their efficiency, as discussed in Section <ref>. 

This study also presents a baseline model built with classical digital signal processing methods, as performance reference of ML-based models. Furthermore, we integrated baseline and ML-based model into ensemble one to gain a optimal compromise between accuracy and efficiency resource consumption? on low-power devices. 






§ CLASSIFICATION BASED ON SPECTROGRAMS



In contrast to the computer vision community, where the transition from frequency decomposition based feature extraction to neural feature extraction on raw image data introduced in 2012 led to a significant increase in predictive performance <cit.>, neural feature extraction on raw audio time series did not contribute to a comparable breakthrough yet. Instead most ML techniques for audio pattern recognition (somewhat surprisingly) were relying on image processing pipelines – sometimes with neural networks trained on ImageNet data. Basically, image classification is performed, using a neural network, analyzing the spectrogram rendering of the audio trace.

For this reason, we initially consider the below two models, based on state-of-the-art neural network architectures.

CNN-Mel –
This model contains two 2D convolutional layers for Mel-spectral inputs as feature extraction and two fully connected layers as classifier. Both convolutional layers utilize 4 filters with a 3 × 3 kernel size and ReLU activation function, followed by 2 × 2 max pooling to reduce spatial dimensions. The output of the classifier are normalized by Softmax function as well.

SqueezeNet-Mel – 
This model is based on SqueezeNet <cit.>, an advanced backbone focused on optimizing the efficiency of computer vision applications by strategically reducing parameters, with comparable performance to AlexNet <cit.>. It leverages a so-called Fire module to achieve higher efficiency compared to standard convolutions with only a slight decrease in accuracy. The Fire module significantly reduces the number of parameters by the squeeze layer (1×1 convolutions), which exceedingly reduces the number of channels before the more parameter-heavy 3×3 convolutions in the expand layer. This reduction in intermediate channels means fewer parameters overall.
We aligned its input shape with the Mel-spectrogram and tailored the output for binary classification. Also offering potential for further compression using quantization techniques, SqueezeNet is thus appealing given memory constraints.





Limitations of Spectrograms on Microcontrollers
Practical experience on low-power microcontroller-based devices has shown however that producing and manipulating mel-spectrograms from audio signal streams on such devices is problematic. It inserts an additional step in the processing pipeline, which increases latency and memory requirements. Prior work (such as <cit.>, section 6.2) concludes that the CPU bottleneck is not just the model inference time, but also the spectrogram calculation, which, compared to inference alone, almost doubles latency. Other previous work such as <cit.> measures on a quite powerful STM32F7 microcontroller that computing and writing in memory the mel-spectrogram of 30 columns and 40 frequency bands takes approximately 1 second.  Note that these spectrogram dimensions are much smaller than our requirements (184 columns × 80 frequency bands) hence even more latency can be expected in our case.

For these reasons, we next explore pipelines which skip the spectral pre-processing stage as described below.



SqueezeNet<cit.> is a CNN characterized by comparable performance to AlexNet while using approximately 50 times fewer parameters. The compact architecture of SqueezeNet is achieved by employing eight unique building blocks called the Fire Modules. Each Fire Module starts with a "squeeze" convolution using 1x1 filters to reduce input channel dimensions, thereby cutting down on computational load. This initial step compresses data into fewer channels. Following this, an "expand" convolution layer mixes 1x1 and 3x3 filters: the former expands channel numbers, while the latter captures spatial details. Moreover, SqueezeNet offers the potential for further compression using quantization techniques, making it an efficient choice given memory constraints.

The implemented SqueezeNet-Time-Series architecture uses a reduced number of feature maps in compare to initial SqueezeNet. This compressed model has in the first convolutional layer 0.3 times the number of filters that initial model has in this layer. In the fire modules in convolutional layers, we use the integer value of 0.3 times the number of squeeze filters given to this layer in initial model. Since the number of expand filters depends on the number of squeeze filters and is four times greater, it also undergoes compression. The structure of the fire module in SqueezeNet-Time-Series is presented in Table 2. Therefore, if the initial number of squeeze filters in the fire module is x, the compressed number of squeeze filters in this fire module will be int(0.3× x). Then both the expand layers will have 4× int(0.3× x) filters.

Additionally, since the initial model uses 2D convolutions and time-series are 1D arrays, we converted all 2D convolutional layers to 1D convolutional layers. The model was trained using time-series data for 10 epochs with the Adam optimizer and cross-entropy loss function. After training, this model was further quantized using different post-training quantization techniques.





For spectrogram approach we aimed to create an ultra-compact CNN-Mel-Spec model, comprising just two convolutional layers and two fully connected layers. The input for the model is a log Mel spectrogram with a size of 184x80. Both convolutional layers utilized 4 filters with a 3x3 kernel size and ReLU activation function, followed by max pooling with a 2x2 pooling window to reduce spatial dimensions. After the last convolutional layer and max pooling, the data were reshaped into a one-dimensional vector for input into the subsequent fully connected layer. The first fully connected layer had 8 neurons and utilized ReLU activation. The final fully connected layer contained two neurons with a linear activation function. For optimization, we employed the Adam optimizer with a learning rate of 0.001, and the model was trained over 15 epochs, using cross-entropy as the loss function. In an effort to reduce model size, we converted it to TFLite format and applied several post-training quantization techniques: float16 quantization, dynamic range quantization, and full integer quantization.













§ CLASSIFICATION BASED ON TIME-SERIES


Contrary to the models described in the previous section and inspired by the success of neural feature extraction in the computer vision domain <cit.>, the pipelines we aim at next take directly the raw audio signal time-series as input. More precisely, we designed the three models described below.

CNN-Time –
This simple model performs feature extraction with two sequential (1D) temporal convolutional layers followed by average pooling. A max pooling is inserted in-between to reduce the dimension of feature maps and enhance the non-linearity of the network. Two fully connected layers act as classifier at the end of the network, with the probability outputs normalized by Softmax function.  

Transformer-Time –
Inspired by <cit.>, we designed an efficient, attention-based model with only one temporal convolutional layer and one single-head transformer. The convolutional and pooling layers serve as feature extraction to transfer raw audio signal into embeddings for the transformer. The activations of the fully connected layer are normalized by Softmax function as well.



SqueezeNet-Time –
This model is based on SqueezeNet <cit.>. We tailored this basis to align the input shape with the audio segment, and narrowed the output channels for binary classification. To adapt the backbone to time-series input, all 2D convolutional layers are replaced by temporal convolutional layers. Furthermore, to have a more compact structure, we decreased the filter number of all convolutional layers by roughly 70%, as presented in Table <ref>.





§ ADDITIONAL MODEL OPTIMIZATIONS FOR TINYML


In this study, we used two optimization techniques to compress the models and further reduce their memory consumption: model quantization on the one hand and on the other hand partial convolution, as described below.







 §.§ Model Quantization

Quantization reduces the model size and inference time by converting the weights and activations from higher precision (e.g., 32-bit floating-point) to lower precision (e.g., 8-bit integer) <cit.>. In this study we adopted PTQ with weights and activations quantized in 8-bit integer. To avoid substantial loss of prediction performance, the training dataset was used to find the optimal scale factor and zero-point of the activations.



 §.§ Partial Convolution

We observed that the outputs of the Conv1D layer require significant memory, which impedes deployment on resource-constrained tiny devices. As presented in Table <ref>, the peak memory consumption occurs at the first Conv1D layer both in CNN-Time and Transformer-Time with 768  and 3 , respectively. This indicates that deployment on resource-constrained tiny devices is impractical.

To address this issue and inspired by <cit.>, 

we exploited the fact that average pooling can be computed iteratively point-by-point on the output channels of the final Conv1D layer over a small sliding window of inputs. Each point in the output channels depends only on a small subset (kernel window) of the outputs from the previous layer. That is, for a block of L Conv1D layers following by an average pooling layer, its output y can be iteratively computed on the input sequence x(n), n= 1 … N as following:

    y_j(k)    = y_j(k-1) + 1/N_L A^L_j(k), 
    
        k    = 1 … N_L, j = 1 … C_L . 
     
        A^l_c(n)    = ∑_i=1^C_l-1 W^l_c(i) · [ A^l-1_i(n-K^l/2) ⋯
       =4A^l-1_i(n) ⋯ A^l-1_i(n+K^l/2) ]^T 
    
        A^0(n)    = x(n) , 
    
        l    = 1 … L, c = 1 … C_l, n = 1 … N_l, N_0 = N

where y_j(N_L) is the result of the average pooling of the j-th channel; A^l_c(n) denotes the n-th point in channel c of the l-th Conv1D layer calculated with kernel weights W. Each Conv1D layer contains C_l channels with filter size of K^l and output size of N_l. Without loss of generality, the formula of partial convolution contains only the linear components and considers only one-dimension case for the sake of simplicity; it can be generalized in high-dimension and integrated with non-linear building blocks (e.g., stride, non-linear activations, pooing layers, etc.).

Figure <ref> provides a comparison between classic convolution and partial convolution. Unlike classic convolution where entire channels (C × N) are computed and stored before being processed by the next layer, partial convolution requires only a small part of the channels (C × K, K << N) for each layer. In our case with K=3 and N=48000, it can be expected roughly 16000 × smaller memory consumption, making the implementation more suitable for tiny devices. We applied this strategy to the Conv1D layers combined with average pooling in the CNN-Time and Transformer-Time architectures.










§ TINYCHIRP PERFORMANCE EVALUATION


In the following, we conducted a comprehensive evaluation of the discriminative performance of TinyChirp with different TinyML models, as well as of resource consumption and computation time on typical low-power boards based on microcontrollers.



 §.§ Classification Performance Comparison







Our evaluation process began with the generation of ROC curves and the corresponding AUC on the training dataset, providing a clear picture of the models' abilities to distinguish between classes. As shown in Figure <ref>, the spectrogram-based models – CNN-Mel and SqueezeNet-Mel – achieved the best classification performance with AUCs of 1.0 and 0.99, respectively, followed by time-series models – CNN-Time and Transformer-Time – both with 0.98 AUC. Unexpectedly, with a more complex structure, SqueezeNet-Time performs worst (lower AUC than the baseline).


We next measured accuracy, precision, recall, F_1- and F_2-score, shown in Figure <ref>. Spectrogram-based models achieve the highest metrics for the whole range of threshold values. For time-series models, Transformer-Time worked overall better than CNN-Time; Again, SqueezeNet-Time performs worst (consistent with the above ROC analysis). 


Next, we focused on finding the threshold settings for optimal  F_2-score prioritizing recall, i.e. reducing the likelihood of mistakenly discarding records of the targeted birds). The results are shown in Table <ref> for training data, thereafter verified on test (previously unseen) data as shown in Table <ref>. Note that this threshold influences energy consumption: more target bird classification leads to more energy consumption (as data needs to we stored on to SD card in this case).

These results are promising and allowing to focus only on the relevant audio segments (10% of the total recording time), as the classification performance is compelling. Thus, TinyChirp can save 90% of SD card space compared to the initial monitoring scenario (Section <ref>), potentially extending the deployment time to 18 weeks. However, complementary experiments must now be carried out on targeted hardware, so as to evaluate other key metrics: computation time, energy consumption and memory footprint on typical microcontroller-based devices. The next sections focus on that part.




presents the optimized setting with corresponding metrics on the training dataset. This approach ensured that our models were fine-tuned to prioritize recall, thereby reducing the likelihood of missing critical positive cases. Also we emphasize that the decision threshold is directly related to energy consumption in the use case considered here: more positive predictions will lead to higher energy consumption. Hence high precision is desirable for low energy consumption, but recall should be traded with caution against energy consumption.

After determining the optimal threshold setting on the training dataset, we validated this threshold using the test dataset, as described in Table <ref>. This step was crucial to ensure that the threshold not only performed well on the training data but also generalized effectively to unseen data. By confirming the optimal threshold on the test dataset, we demonstrated the robustness and reliability of our TinyML models' predictive performance.













 §.§ Performance Evaluation on Microcontrollers


Experimental Setup – For our measurements, we used a common low-power off-the-shelf board: the nRF52840 Development kit (nrf52840dk). This board is based on an ARM Cortex-M4 processor, with a typical memory budget: 1  of Flash memory and 256  of RAM. As embedded software base we used RIOT-ML <cit.>. 

Throughout the experiments, we monitored the (RAM/Flash) resource footprint using built-in diagnostic tools provided by RIOT-ML and RIOT <cit.>, as well as external measurement equipment as described below.

Energy consumption was measured using an ampermeter to gauge the board's energy efficiency. This involved recording the current draw of the MCU and the external storage system (SD card) in different stages. We used a voltage regulator to supply 3.3  DC output and a logger to capture detailed current profiles over time.

Scope of the Measurements – Besides the model inference stage, we also considered the resource footprint of pre-processing during bird song recognition. Pre-processing refers to down-sampling and calculation of Mel-spectrogram as described in Section <ref>.

Also, the energy consumption in MCU's lowest power mode (idle stage) was measured as the reference of absence of sound with sufficient intensity.



Measurement Results on Microcontrollers – The first striking observation is that  SqueezeNet-Time and SqueezeNet-Mel were not deployable due to OOM (exceded memory budget) and were thus excluded in this experiment. All other models fit in the constraints of memory and storage budget, even when discounting the memory footprint of the OS and the network stack for wireless communication. 

The resource footprint of different stages and models measured on an nrf52840dk board are shown in Table <ref>.  We observe that the model inference stage takes the bulk of the overall energy consumption, as it requires the most process time and MCU in full power mode. 
As expected, the idle stage requires the least power, as all peripherals and the MCU are in the lowest power (standby) mode (until woken up by sound with sufficient amplitude). 











Spectrogram vs Time-series – Considering only the prediction performance in Section <ref>, it would seem natural to choose CNN-Mel as the best candidate, since it also provides the best inference latency. However, taking pre-processing time into account (which for CNN-Mel includes producing a spectrogram) shows a different picture. It takes roughly 2.4 seconds in total for a 3 second audio snippet, which is dangerously near the real-time constraint and costs high power consumption. With the best prediction performance and the lowest total compute latency among time-series models, the Transformer-Time now appears like the best choice to deploy on tiny devices.



Next, we consider jointly the overall resource consumption and the classification performance to discuss variants for the decision strategy described in Section <ref>:

Baseline Only –
With the Baseline t_low (Step 1 and Step 3 only) setting, the required computational resources are minimal, while achieving a recall of 0.99; however, 
due to the low precision, many false positives are stored, leading to the lowest storage efficiency.

TinyChirp Skipping Baseline –
With this variant, the TinyML model is applied immediately, without pre-screening by the baseline (Step 2 and Step 3 Only), leading to high predictive performance and high storage efficiency, but requiring higher energy consumption. Even when using the Transformer-Time with the lowest energy footprint, this requires more than 9 times the energy required for the baseline. This variant is ideal for devices currently running out of storage space while still with substantial battery left.

TinyChirp –
This variant refers to the full series of Step 1-3 described in Section <ref>, whereby the samples not discarded by the baseline are double-checked by the TinyML model. This approach offers a good compromise of compute, storage and energy footprint, by discarding early, without impacting too much the overall prediction performance achievable by the TinyML model on its own. 

TinyChirp with Power-saving –
This variant (Step 1-3 with the power-saving flag enabled) provides a more energy-efficient twist to TinyChirp.

However, the overall precision is determined by the threshold setting t_high, which causes an extra 40% of false positives to be stored onto the SD card, leading to lower storage efficiency. This variant seems more suitable for currently running out of battery, but with sill sufficient storage resources.






§ CONCLUSION

In this paper we present an empirical study of a bioacoustic monitoring use case showing that machine learning is usable directly on autonomous recording units based on low-power microcontrollers. We focused on a concrete use-case: monitoring Corn Bunting birds in rural UK. We publish a data set of the recordings, based on which we demonstrated that TinyML can perform on-device high accuracy classification for bird song classification. Compared to the traditional approach used so far, our approach TinyChirp can extend from 2 weeks to 18 weeks – almost a full season – the time intervals until which researchers must harvest overfull data storage embedded on the recording units deployed in the field. Our experiments also observe hands-on how performing inference directly on the audio time-series is the better approach on microcontroller-based hardware, compared to inference on mel-sepctrograms typically used in audio pattern recognition. Beyond the specific use case of monitoring Corn Bunting birds' songs, based on our general-purpose TinyML approach and our open-source code implementation integrated into the RIOT operating system, TinyChirp is applicable to a wide range of similar bioacoustic use cases, using a large variety of low-power microcontroller-based hardware.




§ ACKNOWLEDGMENT

The research leading to these results partly received funding from the MESRI-BMBF German/French cybersecurity program under grant agreements No. ANR-20-CYAL-0005 and 16KIS1395K. We acknowledge the receipt of media from the Macaulay Library at the Cornell Lab of Ornithology and the Xeno-Canto database.  Felix Bießmann received support from the Einstein Center Digital Future, Berlin,
the German Research Foundation (DFG) (528483508 - FIP 12), the Federal Ministry of Economic Affairs and Climate Action (RIWWER 01MD22007H), the German Federal Ministry of the Environment (GCA 67KI2022A) and the German Federal Ministry
of Education and Research (KIP-SDM 16SV8856).


IEEEtran 









 Paper Title: Rapid and Power-Aware Learned Optimization for Modular Receive Beamforming.




An Empirical Study on Challenges of Event Management in Microservice Architectures
    Yongluan Zhou
    August 5, 2024
==================================================================================






§ ABSTRACT  

mimo systems play a key role in wireless communication technologies. 
A widely considered approach to realize scalable mimo systems involves architectures comprised of multiple separate modules, each with its own beamforming capability. Such models accommodate cell-free massive mimo and partially connected hybrid mimo architectures. A core issue with the implementation of modular mimo arises from the need to rapidly set the beampatterns of the modules, while maintaining their power efficiency. This leads to challenging constrained optimization that should be repeatedly solved on each coherence duration. 
In this work, we propose a power-oriented optimization algorithm for beamforming in uplink modular hybrid mimo systems, which learns from data to operate rapidly. We derive our learned optimizer by tackling the rate maximization objective using projected gradient ascent steps with momentum. We then leverage data to tune the hyperparameters of the optimizer, allowing it to operate reliably in a fixed and small number of iterations while completely preserving its interpretable operation. We show how power efficient beamforming can be encouraged by the learned optimizer, via boosting architectures with low-resolution phase shifts and with deactivated analog components.  Numerical results show that our learn-to-optimize method notably reduces the number of iterations and computation latency required to reliably tune modular mimo receivers, and that it allows obtaining desirable balances between power efficient designs and throughput.













§ INTRODUCTION



Massive mimo technology is at the core of future wireless communications, with extremely large implementations expected to enable meeting the constantly growing demands in throughput and coverage <cit.>. However, implementing communications with hundreds or thousands of antenna elements gives rise to several challenges, including issues related to cost, power consumption, deployment, and scalability <cit.>. To overcome these challenges, a leading approach to massive mimo  adopts modular designs. In such cases, an extremely large mimo  array is implemented using many distinct modules, each equipped with a (smaller) antenna array, and possibly possesses local beamforming capabilities. These modular architectures can take different forms, e.g., distributed cell-free massive mimo <cit.>, centralized interconnections of multiple panels <cit.>, or radio stripes <cit.>.

While modular implementations of massive mimo alleviate cost and power issues and support distributed implementation, they also introduce design and algorithmic challenges. In particular, beamforming is not carried out by a single cpu that manipulates the signal at each element; instead, it is divided between the cpu and a set of separate arrays termed panels, that are possibly constrained in their hardware capabilities, and are likely to be subject to strict power constraints. Moreover, the setting of the modular beampattern is done on each coherence duration, necessitating rapid adjustments of the modules <cit.>.

Since modular beamforming is divided between a cpu and multiple panels, it can be viewed as a form of hybrid beamforming <cit.>. Hybrid beamforming divides the processing of the signals between conventional digital processing and an additional constraint antenna interface, typically representing analog hardware <cit.>. Such hybrid mimo is often studied in the context of RF chain reduction <cit.>, most commonly assuming phase shifter based analog processing <cit.>. However, hybrid mimo can  also represent emerging antenna architectures based on true-time-delay <cit.>, vector modulators <cit.>, metasurfaces <cit.>, or leaky waveguides <cit.>. 

The main task of hybrid beamforming can be viewed as the conversion of csi into a beampattern achievable through a divisible architecture <cit.>. Various methods have been proposed for tackling this task (see, e.g., overview in <cit.>). As hybrid beamforming can be represented as a constrained optimization problem, a conventional approach achieves this via iterative optimizers, see, e.g., <cit.>. However, these approaches typically involve lengthy optimization with numerous iterations, which can be challenging to carry out efficiently within the extremely short time frames required for determining beamforming (e.g., on the order of 0.1 milliseconds <cit.>).

To facilitate rapid beamforming setting, recent works proposed using deep learning tools, taking advantage of their ability to learn to tackle challenging optimization problems with fixed and limited latency <cit.>. This can be achieved by training conventional deep learning architectures, such as cnn, to map mimo csi into  analog and digital precoders <cit.>, or by
combining optimization with deep learning for hybrid beamformers <cit.> as a form of model-based deep learning <cit.>. However, while hybrid beamforming can be viewed as encompassing modular beamforming, its associated algorithms are not tailored for exploiting modular structures, and do not account for the limited power considerations imposed in such forms of distributed mimo systems. Moreover, current learning-aided designs often focus on a given configuration, and the optimizer needs to be retrained when using a different array (or adding another module), limiting scalability. Finally, existing works typically focus on downlink beamforming, with the uplink being much less studied.

Uplink modular beamforming was considered in <cit.>. There, the focus was to characterize trade-offs of its decentralization under narrowband communications and sparsity-based power-oriented constraints for loss-less processing, rather than establishing algorithms for tuning modular beamformers. In practice, the regimes of interest are typically those where the distributed and power constrained operation is likely to yield some rate loss compared to centralized fully digital mimo systems. This gives rise to the need to design methods for tuning hybrid modular receive beamformers for such practical and challenging settings.

In this work, we develop a rapid and power-aware tuning algorithm for setting modular receive mimo beamformers. 
Our approach combines principled optimization with machine learning  via model-based deep learning <cit.> to realize  uplink modular beamforming optimizers operating reliably with fixed and low latency. By doing so, we are able to utilize data to optimize the optimizer while maintaining its comprehensibility and transferability across different configurations. 

We focus on modular beamformers implemented using conventional phase shifters, and leverage the interpretable operation of the optimizer to (i) facilitate scalable design, where the same learned optimizer can be utilized with different numbers of modules; and (ii)  encourage power efficient designs for each panel. The latter is achieved by enabling deactivation and low-resolution implementation of the active phase shifters at each module. 
We formulate the modular uplink beamforming task as a constrained maximization of the sum-rate objective in multiple-access channels. We then specialize pga steps with momentum for the considered setting, being a suitable candidate accelerated iterative optimizer for directly optimizing the rate objective, based on which we develop our learning-aided optimizer. 

Our main contributions are summarized as follows:

    
  * First-order optimization for modular receive beamforming:
We formulate the receive modular beamforming task as a constrained optimization, and derive the corresponding first-order iterative optimizer based on pga steps. We specifically characterize the corresponding gradient steps and projection methods, which enable (lengthy) tuning of modular receive beamformers.
    
    
  * Rapid and scalable learned optimization: To enable rapid tuning, we convert the identified pga steps into a machine learning model operating with a fixed and low latency. This is achieved by fixing the number of iterations while using the hyperparameters of pga with momentum as trainable machine learning parameters. We show that the interpretable operation of the  resulting design allows it to be applied to different number of modules (which can be much larger than those used in training), thus facilitating scalability. 
    
    
  * Power-aware learned optimization: We  leverage the fact that the resulting machine learning model is comprised of interpretable pga iterations to incorporate two forms of power-aware considerations: the ability to boost sparse modular architectures with deactivated phase shifters, and the support of low-resolution phase shifters. As the resulting constraints render pga non-differentiable, we approximate its operation for learning purposes via surrogate projection mappings that enable its representation as a machine learning architecture.
    
    
  * Extensive experimentation: We extensively evaluate our learned optimizer for both synthetic channels as well as channels obtained from the physically compliant quadriga model <cit.>. Our numerical studies systematically show that learned optimization enables reliable tuning of modular beamformers with reductions of 10-30 × in iterations compared to conventional optimizers, and that the power-aware design enables efficient operation with only a minor degradation in the accuracy compared to unconstrained designs. 



The rest of this paper is organized as follows: Section <ref> presents the system model of uplink modular beamforming and formulates its design problem. Section <ref> presents our learned optimization method without accounting for power considerations, which are introduced into the method in Section <ref>. Empirical results are reported in Section <ref>, and Section <ref> concludes the paper.


Throughout this paper, we use  boldface lowercase and boldface uppercase letters for column vectors and matrices, respectively, while [M]_i,j is the (i,j)th element of a matrix M.
We use calligraphic fonts for sets, while ℤ and ℂ denote the integer and complex numbers, respectively. The operations (·)^T and (·)^H are used for transpose and conjugate transpose, respectively. I_N refers to the identity matrix of size N. The gradient of the function f(·) with respect to the matrix W is denoted by ∇_Wf(·), while δ_·,· is the Kronecker delta function, and ∡ (·) is the (element-wise) phase operator (in radians).








§ SYSTEM MODEL




In this section we introduce the system model for receive beamforming with modular hybrid mimo receivers. We commence with presenting the considered models for the communication channel and the receiver processing in Subsections <ref>-<ref>, respectively, after which we formulate the beamforming design problem in Subsection <ref>.



 §.§ Communication System



We consider a multi-band multi-user uplink mimo system. The setting involves K single antenna users that communicate with a (possibly distributed) mimo receiver,  over B frequency bins. The transmitted signals are received by M antenna elements, that are divided into multiple panels. The channel output at the bth frequency
is given by

    y[b]=H[b] s[b] + w[b],      b=1,…,B.

In (<ref>), H[b]∈ℂ^M× K is the channel matrix at frequency bin b, while  s[b] ∈ℂ^K represents the transmitted symbols, which are comprised of i.i.d. entries with power ρ_s>0. The signal w[b] ∈ℂ^M is awgn with variance σ_w^2>0.




 §.§ Modular Hybrid Receive Processing



The signal received by the M antenna elements is processed by a centralized cpu. The resulting digitally processed signal is obtained from the channel output in (<ref>) by undergoing modular receive beamforming, where the M antennas are divided into P panels. This processing, illustrated in Fig. <ref>, includes two stages: (i) Panel-wise beamforming, and (ii) Panel-cpu connectivity.



  §.§.§ Panel-Wise Beamforming

The first stage is composed of receive beamforming performed by each panel separately. Each panel has N antennas and L output ports, where the data processing takes place within the pth antenna panel and is described via a N× L frequency-invariant matrix W_p, for each p=1,…,P. The panels are assumed to be hardware limited, which we model using the set of feasible mappings W⊂ℂ^N× L, such that W_p ∈W. 

We consider the following types of analog constraints:

	
  *  Unconstrained Phase Shifters: Here, the analog processing is comprised of phase shifters, where the phase of component can be set arbitrarily, i.e., 
	
    W_ uph = {W∈ℂ^N × L: |[W]_n,l| = 1, ∀ (n,l)}.

	
	
  *   Sparsified Phase Shifters: As phase shifters are active components with non-negligible power consumption, a candidate approach to maintain power-efficient operation at the panels is to allow some of the phase shifters to be turned off <cit.>. This type of feasible set is represented as
	
    W_ sph = {W∈ℂ^N × L: |[W]_n,l| ∈{0,1}, ∀ (n,l) }.

	While the formulation in (<ref>) does not indicate how many phase shifters are inactive, such a requirement can be imposed by constraining the ℓ_0 norm of W to achieve a desired amount of inactive components. 
	

	
	
	
  *   Quantized Phase Shifters: An alternative  approach to limit the power consumption on the panel side is to restrict the phase shifters to support only a discrete set of feasible phases, i.e., support quantized phases shifts <cit.>. Such a constraint with Q-level quantized phases is modeled by the feasible set
	
    W_ qps^Q  = {W∈ℂ^N × L :       ∡[W]_n,l/2π· Q ∈ℤ ,
    	
       |[W]_n,l| = 1,  ∀ (n,l) }.


Note that for all considered feasible sets, the phase shifting components remain frequency-agnostic, i.e., W_p is invariant of the frequency index b. Moreover, we note that  approaches <ref> and <ref>, which aim to improve the power efficiency of the panels, are complementary. Accordingly, one can restrict the analog processing at the panels side to be both sparse and quantized, i.e., have some of the phase shifters deactivated, while the active ones take discrete phase values.


 



  §.§.§ Panel-CPU Connectivity

We model the connectivity to the cpu  via the binary matrix A∈{0,1}^L· P × T, with T being the number of cpu inputs. This matrix is also frequency invariant, and its entry [A]_l · p,t denotes whether the lth output of the pth panel is connected to the tth cpu input. 


The cpu input  z[b] is related to the channel output y[b] via

    z[b]=A^H W^H  y[b],

where W=  blkdiag{W_1,…,W_P}. This constraint comes from the decentralized structure of our communication system, which takes place within each panel. Note that when A is the identity matrix and L=1, this formulation specializes the common form of partially-connected hybrid beamformers <cit.>. 




 



 §.§ Problem Formulation



We aim to design the modular receive beamformer detailed in (<ref>) to maximize the achievable rate. To formulate this, we define the equivalent channel at the bth frequency bin as 


    G≜WA∈ℂ^M× T.

For the uplink setup detailed in Subsection <ref>, the average achievable sum-rate is given by (see, e.g., <cit.>)

    R(G; {H[b]}) = 1/B∑_b=1^B log|I_M +   ρ_s/N_0G(G^HG)^-1
       G^H
        H[b]H^H[b]|.




Using the rate expression in (<ref>), we can formulate the setup of modular receive beamforming as a constrained optimization problem. Since the matrix A represents the hardware connectivity between the panels, it is assumed to be constant. Accordingly, for a given channel realization {H[b]}_b=1^B}, the modular beamforming task is formulated as the recovery of 

    W^⋆   = max_W=  blkdiag{W_1,…,W_P} R(WA;  {H[b]}) 
       subject to W_p ∈W,   ∀ p∈{1,…,P}.


The problem as it is formulated in (<ref>) considers a generic setting for the feasible panel-wise  computation set W. As such, it accommodates both the unstructured setting <ref> as well as the power-oriented <ref>-<ref>. However, when required to operate with sparse panels <ref>, one is likely to also impose an additional constraint on the overall number of active phase shifters  in W. Similarly, imposing <ref> assumes that one operates with a fixed and known resolution dictated by Q.

 

In some cases (e.g., when T ≥max (⌊ MK-L/K+1 ⌋,K), B=1, and while allowing W = ℂ^N × L), one can identify the optimal setting <cit.>. Even in such cases,  tackling (<ref>) typically involves an iterative optimization procedures. In practice, the modular beamformer must be set within a coherence duration (which can be of the order of 0.1 milliseconds <cit.>). To account for this requirement, we limit the number of iterations for these calculations, allowing at most J iterations. To facilitate identifying an optimizer that can reliably tackle (<ref>) within J iterations, we assume access to a data set D composed of |D|>0 past channel realizations, i.e., 
 
    D={{H_i[b]}_b=1^B}_i=1^|D|.



 






§ UNCONSTRAINED UNFOLDED MODULAR BEAMFORMING


In this section, we study modular beamforming design while focusing on unconstrained phase shifters, i.e., constraint <ref>. The  unfolded modular beamforming algorithm, which is geared towards enabling rapid adaptation, serves as a starting point for considering the more restrictive constraints <ref>-<ref> in Section <ref>. In particular, we  first formulate our baseline optimizer of pga with momentum in Subsection <ref>. Then, we convert it into a machine learning model in Subsection <ref>, and provide a discussion in Subsection <ref>.





 §.§ PGA with Momentum



A direct approach to tackle (<ref>) employs first-order methods <cit.>, e.g., pga. This scheme starts from some initial guess W_0, and iteratively takes gradient steps towards the maximization of (<ref>), projected to meet the constraints. pga relies on the ability to compute the gradient of R(WA;  {H[b]}) with respect to W, denoted ∇_WR(W,A). The gradient of the rate is stated in Theorem <ref>.

The gradient of (<ref>) with respect to W is 
    
    
    ∇_W    R(W,A) =
            1/B∑_b=1^B
            (ρ_s/σ_w^2A((G^HG)^-1G^H 
       ×H[b]H[b]^H)
            (I +ρ_s/σ_w^2(G(G^HG)^-1G^H 
       ×H[b]H[b]^H))^-1(I-G    (G^HG)^-1G^H)) ^H.




    The proof is given in Appendix <ref>. 


In addition to the gradient in (<ref>), one should also be able to project onto the feasible constraints in (<ref>)  to formulate  pga.
For the non-power-constrained feasible set (<ref>), the projection operator for a matrix W based on <ref> is


    [P_W_ uph(W)]_i,j   = [W]_i,j/|[W]_i,j|·δ_⌊i/N⌋, ⌊j/L⌋.

In (<ref>), the element-wise operation divides each entry of the matrix W by its own magnitude, and then sets it to zero if it falls out of the main block diagonal. This effectively keeps each panel-wise processing made up of phase-shifters only.


Although first-order methods often require a large number of iterations to converge, they can typically be accelerated by introducing momentum. Momentum is also useful for dealing with non-convex objectives, as in (<ref>), since it can prevent getting trapped in local minima and saddle points. The resulting pga with momentum is summarized as Algorithm <ref>.










 §.§ Unfolded PGA with Momentum


Algorithm <ref> allows tackling (<ref>) while meeting the requirement to operate with J iterations. However, iterative optimizers such as pga are designed to tackle convex optimization problems and to operate with asymptotically large number of iterations. There, convergence can be guaranteed with fixed  hyperparameters {μ,β} (i.e., independent of the iteration index j), which can be tuned manually. In Algorithm <ref>, we apply this convex optimizer to a non-convex problem (<ref>) while operating with a fixed and small number of iterations J. In such cases, the  setting of the hyperparameters {μ_j,β_j} greatly affects performance <cit.>. This motivates leveraging the available  data in (<ref>) 
to tune these hyperparameters, by converting Algorithm <ref> into a discriminative machine learning model <cit.>, and proposing a dedicated learning mechanism. We next detail the resulting machine learning architecture, its proposed training procedure, and how it can be scaled to large modular settings.




  §.§.§ Architecture



Algorithm <ref> operates with exactly J iterations. It can thus be readily converted into a machine learning model via deep unfolding for learning hyperparameters <cit.>. Here, each iteration is treated as a layer in a multi-layer architecture, whose trainable weights are the hyperparameters. 

We increase the abstractness of the trainable architecture by allowing it to use different values of μ and β not just per iteration, but also per entry in W. Accordingly, the scalar values  {μ_j, β_j} that multiply matrix quantities in Algorithm <ref> are replaced with block-diagonal matrices α_j, β_j that multiply the matrix quantities in Algorithm <ref> element-wise.
The resulting architecture obtained from Algorithm <ref> is illustrated in Fig. <ref>, and has 2· J · P · N · L trainable parameters,  written as
 
    θ=  {α_j, β_j}_j=0^J-1.







  §.§.§ Training



Converting Algorithm <ref> into a machine learning model facilitates tuning θ  in (<ref>),
which are the hyperparameters of Algorithm <ref>. This is achieved using the dataset D in (<ref>), containing past channel realizations taken from the same communication system. Letting W_J(θ) be the modular beamforming matrix computed by Algorithm <ref> with parameters θ, we seek the hyperparameters 

    θ^⋆ = max_θ∑_{H[b]}∈D R(W_J(θ)A;  {H[b]}).


The objective in (<ref>) evaluates θ  based on the output of the unfolded architecture, i.e., W_J. To facilitate the learning of θ, we exploit the interpretability of the architecture, and particularly the fact that each layer should output a gradually improved modular beamformer. Accordingly, we propose a loss measure for optimizing θ that accounts for the output of the intermediate iterations (layers) in a monotonically growing fashion. Specifically, we use a loss function with logarithmic contribution growth between iterations <cit.>, given by

    ℒ_𝒟(θ) = -∑_{H[b]}∈D∑_j=1^Jlog(1+j)/|D| R(W_j(θ)A;  {H[b]}).

The resulting training algorithm using mini-batch stochastic gradient descent is stated as Algorithm <ref>.
The negative sign in the loss  is added because our objective is to maximize the rate, rather than minimize it.



  














  §.§.§ Transferability Across Configurations



Modular beamforming architectures are scalable,  stemming from the ability to achieve extremely large mimo configurations by adding panels. A core gain of our interpretable learning-aided optimizers is the fact that they support such scalability. Specifically,  what we learn are hyperparameters of an optimization algorithm, which can be applied to different configurations. This indicates that our unfolded optimizer shares this transferability across configurations, thus inherently supporting large and adaptive scale deployments.

In particular, the architecture and training procedure detailed above are formulated for a given modular configuration. Nonetheless, a learned optimizer is transferable, and a learned optimizer trained for a small configuration can be scaled to applied to much larger ones. Specific scalable parameters of interest are the numbers of users K and of panels P:

    
  * Scaling K: we note that the learned hyperparameters in (<ref>) are invariant of the number of users. Thus, a learned optimizer trained for a small number of users can be readily applied to a setting with a large number of users.
    
  * Scaling P: While the setting of the hyperparameters in (<ref>) depends on the number of hyperparameters, as we allow panel-specific hyperparameters, one can scale by reusing hyperparameters across panels <cit.>. For instance, a learned optimizer trained to set {α_j, β_j} for a configuration with P panels can be applied to a setting with c· P a positive integer c, by e.g., setting its hyperparameters to be {I_c ⊗α_j, I_c ⊗β_j }, where ⊗ denotes the Kronecker product, or by assigning the added panels an average of the learned hyperparameters.

The above approaches to transfer a learned optimizer across different configurations are empirically demonstrated to yield effective and scalable optimization in Subsection <ref>.















 §.§ Discussion



The proposed unfolded pga  is designed to learn from data how to traverse the optimization loss surface, allowing to achieve improved sum-rate within J iterations, as shown in Section <ref>.
The resulting learned θ, which represents the per-iteration hyperparameters used for setting the modular beamformer, may exhibit certain counter-intuitive characteristics, such as negative gradient steps, as means of coping with non-convexity. As the resulting algorithmis a learned implementation of a principled optimizer, we share its scalability, which is highly desirable in modular architectures. 

While our formulation considers panels comprised phase-shifters, one can consider alternative hardware, such  as distributed active metasurfaces <cit.> and phase-time arrays <cit.>. 




The fixed and limited number of iterations impacts the computational complexity of setting the modular beamformers for a given channel realization. To quantify this complexity, we examine a single iteration of Algorithm <ref>. The computational burden is mainly dominated by the gradient computation step, formulated in (<ref>). This computation has a complexity order (in complex multiplications) of

    𝒪(B(T^3 + M^2(T + K)) + M (T^2 + LPT + TK)).

The remaining steps in the algorithm include the update step and the projection step (<ref>). The complexity order of these steps is  𝒪(MLP) for both, which is negligible compared to the gradient computation step. Considering that optimization procedure  operates over B frequency bins, the overall complexity of the PGA+M algorithm with a predefined J iterations, as implemented by the proposed learned optimizer, is of the order of

    𝒞_PGA+M = 𝒪(  JB(   T^3 + M^2(T + K) 
     
        + M (T^2 + LPT + TK))).


The asymptotic complexity characterization in (<ref>) is comprised of multiple summands, where in general none of the terms can be consistently ignored. Nonetheless, in an expected massive modular mimo setting, where the overall number of antenna elements is larger than both the number of users (K) and the number of cpu inputs (T), the dominating term is the one growing quadratically with M, and thus  𝒞_PGA+M≈𝒪(JBTM^2). 
When delving further into the analysis, it becomes evident that the dominant terms (in complex products) arise from matrix multiplication operations.
However, these computations can often be parallelized to reduce computational complexity, making the complexity of matrix inversion, represented by the 𝒪(JBT^3) term, the dominant factor. While it is possible to further reduce the number of multiplication operations through techniques such as sparse and low-rank approximations of the gradients, we leave this aspect to future research.








§ POWER-AWARE UNFOLDED MODULAR BEAMFORMING




Our analysis in Section <ref> focused on the unconstrained setting <ref>.
In this section, we use the unconstrained design for maximizing the rate (<ref>) as a starting point, on top of which we develop a power-aware design. Specifically, we leverage the interpretability of the solution described in Section <ref> to incorporate constraints <ref>-<ref>. We focus on conventional (non-learned) gradient-based optimization subject to  <ref>-<ref> in Subsections <ref> and <ref>, respectively. Then, we  formulate  pga with momentum and the unfolded algorithms tailored for these designs in Subsections <ref> and <ref>, respectively, and conclude with a discussion in Subsection <ref>.





 §.§ Sparsified Modular Beamformer




As detailed in Section <ref>, power efficiency can be enhanced by turning off some of the phase shifters.
Supporting deactivated phase shifters can be mathematically accommodated by constraint <ref>, i.e., as promoting  sparsity of the matrix W. 
To meet this constraint, we rewrite our constrained optimization from (<ref>) as 


    W^⋆   = max_W=  blkdiag{W_1,…,W_P} R(WA;  {H[b]}) 
       subject to  ||W||_0 ≤ S_max
                 W_p ∈W_ sph,   ∀ p∈{1,…,P},

where S_max≤ N· L · P represents  the maximal allowed number of active phase shifters. 
In (<ref>), each phase shifter can be either active or turned off (by (<ref>)). The ℓ_0 constraint guarantees that at most S_max phase shifters are active.

Due to the inherent challenges associated with gradient-based optimization under ℓ_0 constraints, we adopt the common approach of relaxing it into an ℓ_1 constraint (see, e.g., <cit.>). The relaxed optimization problem is 

    W^⋆   = max_W=  blkdiag{W_1,…,W_P} R(WA;  {H[b]})  - λ ||W||_1 
       subject to W_p ∈W_ sph,   ∀ p∈{1,…,P}.

In (<ref>),  λ>0 is an objective regularization coefficient <cit.> whose purpose is to encourage the solution to hold the sparsity constraint in (<ref>). Setting the exact value of such hyperparameters to meet a desired sparsity level S_max is typically elusive, and involves repeating numerical trials to identify the exact value that strikes a desired trade-off between performance (i.e., rate) and sparsity. Nonetheless, once λ is fixed, the gradient of the relaxed objective in (<ref>) can be computed as 

    ∇_W(R(WA;   {H[b]}) -  λ||W||_1) =
       ∇_W R(WA; {H[b]}) - λ· e^j∡W.

where ∇_W R(·) is given in Theorem <ref>.







Relaxing (<ref>) into (<ref>) enables computing the objective gradients via  (<ref>). In order to enable pga-based optimization, one should also define the corresponding projection on the feasible set 𝒲_ sph. To that end, we adopt a flexible extension of the projection operator in (<ref>) to account for deactivated phase shifters, by introducing a magnitude threshold ζ>0. The resulting projection step is given by 

    [𝒫_𝒲_ sph(W)]_i,j=
    [W]_i,j/|[W]_i,j|·δ_⌊i/N⌋, ⌊j/L⌋    |[W]_i,j| ≥ζ, 
    
    0     |[W]_i,j| < ζ.




Projection via (<ref>) is a crucial  step in the tackling the relaxation of (<ref>), ensuring that we set some elements of W to 0 and constrain their magnitude to be either 0 or 1, as described in <ref>.
At the same time, this step preserves the block diagonal structure of W. A candidate choice of ζ is ζ = 0.5; Setting ζ = 0 makes this projection step coincide with the projection step in (<ref>).





 §.§ Quantized Modular Beamformer



A complementary approach to mitigate power consumption of analog modular processing is by discretizing the phase resolution with a few bit representation. Incorporating the resulting constraint <ref> into the rate maximization objective (<ref>), yields the resulting constrained optimization formulation

    W^⋆   = max_W=  blkdiag{W_1,…,W_P} R(WA;  {H[b]}) 
       subject to W_p ∈W^Q_ qps,   ∀ p∈{1,…,P}.

In (<ref>), the set of feasible phase shifts is restricted to take Q quantization levels, using the definition of W^Q_ qps in (<ref>). 

We note that the objective in (<ref>) is the same as in the unconstrained case (<ref>), and thus its gradient can be obtained from Theorem <ref>. Consequently, to formulate the suitable pga steps for tackling (<ref>), one must specify the quantization-aware projection operator. A direct approach to guarantee discrete phase components is by setting each phase element into its nearest discrete value, resulting in 

    [𝒫_𝒲_ qps^Q(W)]_i,j   = e^j ϕ_i,j^⋆·δ_⌊i/N⌋, ⌊j/L⌋
    where ϕ_i,j ^ ⋆   = min_ϕ/2π· Q ∈ℤ 
    |ϕ - ∡[W]_i,j|.

The projection operator in (<ref>) can be viewed as uniform scalar quantization of the phase of each element of W. 





 §.§ Power Aware PGA





Here we present a baseline pga algorithm for the updated set of constraints, as described in <ref> and <ref>. 
This pga algorithm is an extension of Algorithm <ref>, adapted to the power aware objectives described in Subsections <ref> and <ref>.

The resulting formulation is summarized
as Algorithm <ref>. There, we integrate the power-aware gradient formulation (<ref>) with the projection steps (<ref>) and (<ref>) to realize pga.



Algortihm <ref> is designed to tackle the maximization problems outlined in (<ref>) and (<ref>) by proper settings of the projection operator and the regularization coefficient. For instance, by setting the projection to be given in (<ref>) and setting λ = 0, Algorithm <ref> tackles problem (<ref>). Similarly, by using λ > 0 and setting the projection to be (<ref>), it tackles problem (<ref>).

 







 §.§ Unfolded Power Aware PGA



In Subsection <ref>, we showed how data can be leveraged to tune the hyperparameters of Algorithm <ref> to enhance its performance within J iterations, by converting it into a machine learning architecture.

Here, we follow this methodology to convert the power-aware Algorithm <ref> into a trainable architecture. As opposed to the unconstrained case, where the learned hyperparameters were solely associated with the solver, i.e., with pga, in the following we also leverage data to tune parameters associated with the optimization objective <cit.>, and particularly the (typically elusive) regularization coefficient λ. 




  §.§.§ Architecture

Following the conversion of pga with momentum into a machine learning model  in Subsection <ref>, we view the iterative algorithm as a multi-layer architecture. This allows using the  data (<ref>) to learn  hyperparameters that yield desirable settings within a few iterations. 

 Algorithm <ref>  introduces another hyperparameter, λ, associated with optimization objective, and particularly the ℓ_1 norm reduction. Here, we increase the abstractness of the trainable architecture by allowing λ to differ from one iteration to another, i.e., have each iteration taking its gradient step over a different loss surface. Furthermore, we turn the scalar regularization coefficient into a matrix with different values for each entry in W. We add these learnable matrices {λ_j}_j=0^J-1 to the set of trainable parameters θ. Accordingly, we now have  3 · J · P · N · L trainable parameters, written as

    θ=  {α_j, β_j, λ_j}_j=0^J-1.


The power aware architecture is visualized in Fig. <ref> for a given learned hyperparameters θ. For clarity of presentation, in Fig. <ref> we separate the gradient step in Algorithm <ref> into two separate steps, each one dependent on a different hyperparameter.










  §.§.§ Projection

A key issue in deep unfolding is maintaining the  forward path differentiable. This allows the optimizer to be represented as a  trainable architecture that can be tuned with deep learning techniques based on  backpropagation <cit.>. Algorithm <ref> mainly differs from Algorithm <ref> in the projection step. The projection step in both (<ref>) and (<ref>) contains continuous-to-discrete mappings that nullify gradients. 


To overcome this, we distinguish between the projection step utilized when applying the optimizer, and a  differentiable surrogate utilized for evaluating gradients during training. Specifically, we use the following surrogate projections:

    
  * Approximating 𝒫_𝒲_ sph: The projection onto 𝒲_ sph can be viewed as a step function. Accordingly, we approximate it during training using a harp and shifted sigmoid function <cit.> denoted 𝒫̂_𝒲_ sph, given by
    
    [𝒫_𝒲_ sph]_i,j = |W_i,j| = σ (s_m (|W_i,j| - ζ_m)).

    In (<ref>) σ denoting the sigmoid function, s_m denotes a sharpness hyperparameter, and ζ_m stands for the shift, imitating the threshold ζ from (<ref>).

    
  * Approximating 𝒫_𝒲^Q_ qph: For phase approximation with Q  quantization levels, we note that the projection operator is a piecewise-constant monotonically increasing function. Following <cit.>, we  approximate the quantization levels with sum of Q differently shifted sigmoid functions. The resulting phase projection sets ϕ_i,k^⋆ in (<ref>) to be 

    ϕ̂_i,k^⋆= 2π/Q∑_q=1^Q σ(s_p(∡W_i,j - 2π (q-0.5)/Q + π)),

where s_p is a sharpness hyperparameter.



This above formulations  approximate the  projection operators in a differentiable manner, ensuring smooth gradients during training while maintaining the characteristics of the original projection step.
 The ability of (<ref>)-(<ref>)  to approximate their corresponding hard (non-differentiable) projection operators is illustrated in Fig. <ref>.















  §.§.§ Loss Function


We assume access to a dataset D, which comprises past channel realizations.
This data is leveraged to learn the hyperparameters for solving the following constrained optimization problem:


    θ^⋆ =  max_θ∑_{H[b]}∈D R(W_J(θ)A; {H[b]}) 
       subject to: W_J=  blkdiag{W_J_1,…,W_J_P}
       (i)     ||W_J||_0 ≤ S_max,   W_J_p∈W_ sph,   ∀ p∈{1,…,P}. 
       (ii)   W_J_p∈W_ qps,   ∀ p∈{1,…,P}.

 
The constraints on W are either  (i) or (ii), depending on the solved problem, (<ref>) or in (<ref>), respectively. 


The learning objective in (<ref>) under (i) encapsulates a key difference from  the unconstrained one  (<ref>).  Specifically, for sparse phase shifters, the learning objective stems from (<ref>), namely, maximizing R under the ℓ_0 norm constraint. When formulating the resulting pga iterations, we relaxed ℓ_0 norm to the ℓ_1 norm, using the coefficient λ to balance between rate maximization (R) and  sparsification. To translate this into a suitable loss for training that encourages a desired ℓ_1 norm, we introduce the following training loss 


    ℒ_𝒟(θ) =  -1/|D|∑_{H[b]}∈D   ∑_j=1^Jlog(1+j) R(W_j(θ)A;  {H[b]}) 
       - γ(W_J(θ)_1 - S_max)^2,

where γ≥ 0 is a regularization hyperparameter. The loss (<ref>)  is also used with quantized phase shifters, i.e., under (ii), by setting γ=0, as constraint (ii) is guaranteed to hold when using the projection operator 𝒫_𝒲^Q_ qph in modular beamforming. 

The updated loss function (<ref>) is an extension of the unconstrained loss function (<ref>), adding the squared error between the ℓ_1 norm of W_J and S_max. 
The rationale stems  from our desire to encourage W_J_1 (being a differentiable approximation of W_J_0) to approach S_max without having to tune different regularization coefficients for different S_max. This allows training via, e.g., mini-batch stochastic gradient descent, as in Algorithm <ref>.












 §.§ Discussion



The design of the constrained modular beamformer via Algorithm <ref> is based on the pga steps designed for the unconstrained setting in Algorithm <ref>, leveraging its interpretable operation by introducing dedicated projections. Despite the similarity in inference, the procedure involved in learning the hyperparameters from data substantially differs from the unconstrained case, necessitating differentiable approximations of the projection operators during learning, altering gradient computation, and a dedicated training loss. This allows learning to rapidly set constrained modular beamformers with only a minor degradation in the rate, as shown in Section <ref>.

Once the hyperparameters are learned through the unfolded learning procedure detailed in Subsection <ref>, the setting of the modular beamformer in Algorithm <ref> follows the unconstrained pga in Algorithm <ref>. The key differences are the additional phase term in the gradient computations via (<ref>), and the modification of the projection operators via (<ref>) or (<ref>). As none of these operations dominates the computational burden, the complexity of the power-aware design is similar to that of the unconstrained case, analyzed in Subsection <ref>. 

Our power-aware design boosts power efficiency by supporting deactivated and low-resolution phase shifters. The translation of these properties into  power savings is highly dependent on the hardware. We opt for a parametric representation of the desired sparsity level and phase shifter resolution, allowing the unfolded optimizer to be useful for different implementations and different power levels by setting of the corresponding parameters, e.g., S_max and Q. While our learning procedure utilizes the projection surrogates visualized in Fig. <ref>, one can consider alternative differentiable approximations of such continuous-to-discrete mappings, see, e.g., <cit.>. Moreover, additional methods to boost power efficiency by, e.g., deactivating panels rather than phase shifters, can be obtained as extensions to our framework, which are left for future study.




















§ EXPERIMENTAL STUDY



In this section, we numerically evaluate the proposed unfolded algorithm for rapid tuning of uplink modular beamforming[The source code and hyperparameters are available at <https://github.com/levyohad/Power-Aware-Deep-Unfolding-for-Beamforming>]. We describe the experimental setup in Subsection <ref>, while our results for unconstrained and power-aware settings are reported in Subsections <ref>-<ref>, respectively.



 §.§ Experimental Setup


We simulate a communication system, where the channel matrices {H[b]}_b=1^B are obtained from a Rayleigh fading distribution, as well as physically compliant realizations obtained from the quadriga model <cit.>. We treat the number of iterations J as a pre-defined requirement, and train our algorithm to minimize its loss within this number of iterations.

The snr is defined as  1/σ_w^2.


We compare our unfolded optimizer (U-PGA+M) with the following benchmarks: 
 
    
  * PGA + M: running pga with momentum (Algorithm <ref>) with fixed hyperparameters {μ,β}, ∀ j =1 … J, chosen from a discrete grid, to maximize the average R after J=500 iterations for the entire training set. 

    
  * Line Search:   pga where the hyperparameters {μ_j,β_j}_j=0^J-1 are selected for each iteration through grid search to maximize the achievable rate R at each step <cit.>. This method performs an exhaustive search to find the optimal step size for each iteration, independently for each data sample, at the cost of excessive additional processing latency. 
    
    
    
    
  * Manifold Optimization (MO): The optimizer proposed in <cit.> for centralized architectures, adapted to support   modular beamformers. It suggests a relaxed version of the maximization problem (<ref>), taking gradient steps of W toward the fully digital beamformer G_ opt, trying to minimize the error ||WA - G_ opt||^2, and projecting the resulted matrix W on the feasible set, using P_W(·). As in the Line Search benchmark, it performs an exhaustive line search to choose the best step size for each iteration. 
   
    
  * cnn: A cnn based on the architecture of <cit.>, trained to map channel realizations {H[b]}_b=1^B into the phases of W that maximize the channel-rate R. This approach is extremely parameterized. 
    For example, in a system with K=7 users, M=8 antennas, and B=4 frequency bins, the  cnn  consists of 274088 learned parameters, whereas U-PGA+M has only 960 parameters.

 



 §.§ Unconstrained Modular Beamformers


We start the experimental study by considering unconstrained modular beamformers, as described in Subsection <ref>. Our evaluation first focuses on small scale mimo settings, after which we assess transferability to larger systems.



  §.§.§ Small-Scale Hybrid Modular mimo

We first examine two scenarios representing relatively small-scale settings. In Scenario 1, we set (T=5, P=2, L=4, N=20, B=2, K=20), and the channel matrices obey Rayleigh fading. Scenario 2 uses quadriga channels with (T=5, P=4, L=2, N=3, B=4, K=5).  

We use these scenarios to evaluate the ability of our unfolded design to improve the sum-rate achieved within J=10 iterations. To that aim, we depict in Figs. <ref>-<ref> the average sum-rate vs. the iteration index j for Scenarios 1 and 2, respectively.





For fair comparison, we set all optimization-based algorithms to have J=10 iterations, with the cnn having 10 layers, and compare the sum-rate to that achieved by PGA+M  with J=500 iterations. We consistently observe that in the considered scenarios with a limited number of iterations, our unfolded pga outperforms all benchmarks, including the costly and greedy line search hyperparameter setting, and the highly parameterized cnn benchmark. The latter particularly struggles in Rayleigh channels, as  cnns tends to capture spatial relationships, and such relationships do not exist under Rayleigh fading. Similarly, the MO optimizer exhibits slow convergence due to the additional modular structure imposed in our setting of modular beamformers. 


The results after  J iterations illustrated in Figs. <ref>-<ref> are also reported in Table <ref>, along with  latency of the various methods. The latter is evaluated for all algorithms on the same platform (Google Colab premium). 

We notice in Table <ref> that the unfolded algorithm is 30-50 times faster than optimizers that employ line search or run a large number of iterations, while achieving a comparable sum-rate.  
The cnn benchmark has the least latency here, owing to the efficient implementation and natural support of such architectures in platforms such as Google Colab.
This efficiency persists across multiple scenarios, emphasizing our ability to provide reliable modular beamformers with minimal latency, that can be further enhanced by employing dedicated hardware accelerators.















Figs. <ref> and <ref> show the resulted sum-rate for each benchmark in a wide range of snr values, for Scenarios 1 and 2, respectively. 
There, we observe that the gains obtained by the unfolded algorithm while using the same number of iterations, are translated into snr gains of 1.25 - 1.5 dB compared to the pga + M algorithm with the same J iterations.




  §.§.§ Transferable Modular Structures Simulation Results


As explained in Subsection <ref>, one of the key advantages of our learn-to-optimize modular beamforming design is the transferability of learned parameters when scaling the system. Modular structures allow us to train the system with specific settings and then scale it up to a larger number of users or panels without needing to retrain the model. 

Scaling the number of users:
To demonstrate this, we first evaluate U-PGA+M that was initially trained with a small number of users, K=5, and then deployed in a system with a large number of users, K=50, using the same hyperparameters.  This example illustrates the transferability of the unfolded model.
For this evaluation, we introduce Scenario 3: (T=5, P=8, N=4, B=2, K=50) and Scenario 4: (T=5, P=8, N=4, B=2, K=5). In both scenarios the channel realizations are sampled from a Rayleigh fading distribution.
Fig. <ref>, presents the sum-rates achieved by the various benchmarks  on Scenario 3, along with the performance of the U-PGA+M algorithm, which was trained on Scenario 4. These results demonstrate the flexibility of the unfolded algorithm and underscore its advantages in adapting to rapidly changing environments. The average rates at snr of 0  are also reported in Table <ref>.




Scaling the number of panels:
Next, we provide an example of scaling up the number of panels. In this example, we use c = 8 as the scaling coefficient, and increase the number of panels from P to c· P, using {α_j, β_j } = {I_c ⊗α_j, I_c ⊗β_j}.
To this end, we introduce Scenario 5: (T=5, P=2, N=4, B=2. K=7) and Scenario 6: (T=5 · 8, P=2 · 8, N=4, B=2, K=7); both  obey Rayleigh fading.




Fig. <ref>
illustrates the results of various benchmarks trained on Scenario 6, alongside the performance of the U-PGA+M algorithm that was trained on Scenario 5. These results highlight the flexibility of our design, demonstrating how it can be initially trained on a smaller network, perform straightforward calculations, and then be effectively scaled to much larger and more complex systems.





 §.§ Power-Aware Modular Beamformers


We proceed to evaluating the ability of our unfolded optimizer to encourage power-efficient modular beamformers. We divide our evaluation to boosting efficiency via deactivating phase shifters, and via using low resolution phase shifters.



  §.§.§ Sparse Modular Beamformers

We evaluate the effect of designing modular beamformers with deactivated phase shifters by boosting sparsity in W. We focus on Scenario 2,  and apply U-PGA+M using Algorithm <ref>  utilizing the projection step  P_W(·) in  (<ref>). We set the parameter S_max to be 0.75 · N · L · P. 
As the benchmarks detailed in Subsection <ref> do not naturally account for such constraints, we compare here U-PGA+M to unconstrained settings that are based on pga, possibly with line search and/or large number of iterations. 





The resulting rates vs. iteration compared to those achieved without such power-aware design are reported in Fig. <ref>. We observe that by applying Algorithm <ref>, one achieves a 25 % reduction in active components while maintaining a comparable sum-rate and identical running time to unconstrained settings. 
The unfolded algorithm learns to rapidly compensates for the inactive components by tuning the active phase shifters.
Fig. <ref> further highlights the advantages of the unfolded power-aware algorithm by evaluating the resulting sum-rate over different snr. There, we observe that our power-aware design not only allows accurate operation with 25% of the phase shifters deactivated, but it also offers approximately 1 dB improvement in snr over the pga + M benchmark  with the same number of iterations.













  §.§.§ Low-Resolution Modular Beamformers

We conclude our numerical study by evaluating the effect of using quantized phase shifters, and the ability of our power-aware design to account for such constraint when tuning modular beamformers. 
We consider the power-aware  Algorithm <ref>, utilizing the projection step for P_W(·) as specified in (<ref>), when the phase shifters can be configured to merely 4 bits in phase, i.e.,  Q=16. 
The benchmarks detailed in Subsection <ref> are adapted to produce low-resolution phase shifters by altering the projection step (for optimization methods) or quantizing the produced phases (for the cnn).  

The resulting sum-rate vs. iteration and vs. snr for Scenario 1 are illustrated in Figs. <ref> and <ref>, respectively.
The results demonstrate that the proposed unfolded algorithm, which is particularly geared to accounting for low-resolution phase shifters, significantly outperforms other benchmarks, including both slow iterative methods that are not trained as machine learning models, and suffer more notable degradation due to quantization, as well as the  heavily parameterized cnn.

Specifically, while pga + M struggles to overcome the quantization constraints using gradient steps, and tends to converge to a local maximum, the unfolded algorithm effectively learns to traverse the loss surface and projection constraints, achieving a better sum-rate. As shown in Fig. <ref>, the unfolded algorithm consistently performs better across various snr values compared to all other benchmarks.






















§ CONCLUSIONS




In this work, we investigated rapid beamforming design for uplink modular mimo systems under three different sets of constraints: unconstrained phase shifters, sparse phase shifters, and quantized phase shifters. Our proposed algorithm, a momentum-aided gradient-based optimization of the sum-rate objective, boosts power efficient settings while operating  with a fixed and small number of iterations. By transforming the optimizer into a lightweight, trainable model, we harness data to fine-tune its hyperparameters, enhancing its adaptability to various system conditions. Our experimental results demonstrate a significant improvement in the achieved sum-rate compared to benchmarks operating within the same computational-time constraints, while also surpassing other benchmarks in hardware-aware and power-aware designs. 








lemmasubsection 
corollarysubsection 
remarksubsection 
equationsubsection	



 §.§ Proof of Theorem <ref>


For ease of presentation, we define the following auxiliary variables: 
G^†≜ (G^HG)^-1G^H,
P≜GG^†, 
Z[b] ≜PH[b]H^H[b],
and Q[b] ≜I_k + ρ_s/σ_w^2Z[b].
Using these symbols, the objective with respect to which the derivative is taken can be written as

  R(G; {H[b]}) =  1/B∑_b=1^B log|Q[b]|.    


Plugging-in auxiliary variables, we aim at proving Theorem <ref> by  showing that 


    ∇_W R(W,A) = 
           1/B∑_b=1^B
        (ρ_s/σ_w^2AG^†H[b]H[b]^H
        Q^-1
        (I-P)) ^H.

The representation in (<ref>) can be derived using the differential method for matrices as described in <cit.>.
In the following proof the sign : denotes the Frobenius inner product, and the derivations are done using chain rule. For brevity we include our derivation for a single frequency bin.


Since R =log|Q| when B=1, it holds that the differential can be expressed as

    dR =   (Q^-1)^H  :  dQ
         (a)=(Q^-1)^H  : ρ_s/σ_w^2 dZ
    (b)=   (ρ_s/σ_w^2Q^-1)^H  :  dZ
         (c)=(ρ_s/σ_w^2Q^-1)^H  :  dP HH^H
    
         (d)=   (ρ_s/σ_w^2 HH^HQ^-1)^H  : (I - P)   dG G^†

where (a) stems from d Q = 0d I_K + d ρ_s/σ_w^2Z = ρ_s/σ_w^2 d Z, 
(b) and (d) are based on the Hermitian property  inner product, and (c) stems from d Z = d PHH^H, as H is independent of W.

Next, we substitute  P, G^†, and G using chain rule:


    dR =    (ρ_s/σ_w^2  (HH^H)Q^-1)^H  : (I - P)   dG G^†.


Equation (<ref>) can be explained by recalling that P≜G (G^H G)^-1G^H, and thus 

    dP   = dG (G^H G)^-1G^H + G d(G^H G)^-1G^H  
                + G(G^H G)^-10dG^H
    
            = dG G^† - G (G^H G)^-1 (0dG^HG + G^H dG) G^†
       = dG G^† - G (G^H G)^-1G^H dG G^†
       = (I - P)   dG G^†.

In (<ref>), the term dG^H vanishes since G^H is independent of G, and thus independent of W.
Once (<ref>) is achieved, the rest of the proof becomes straight forward, using the chain rule until (<ref>) is reached, i.e., 

    dR =   (ρ_s/σ_w^2 G^†HH^HQ^-1 (I - P))^H  :   dG
     
         =   (ρ_s/σ_w^2 G^†HH^HQ^-1 (I - P))^H  :   dWA
     
        =   (ρ_s/σ_w^2 AG^†HH^HQ^-1 (I - P))^H  :   dW ,

which concludes that ∇_W R(W,A) = (ρ_s/σ_w^2AG^†HH^HQ^-1(I-P)) ^H, thus proving the theorem.


 



IEEEtran




