Provably Efficient Causal Model-Based Reinforcement Learning for Systematic Generalization

In the sequential decision making setting, an agent aims to achieve systematic generalization over a large, possibly infinite, set of environments. Such environments are modeled as discrete Markov decision processes with both states and actions represented through a feature vector. The underlying structure of the environments allows the transition dynamics to be factored into two components: one that is environment-specific and another one that is shared. Consider a set of environments that share the laws of motion as an illustrative example. In this setting, the agent can take a finite amount of reward-free interactions from a subset of these environments. The agent then must be able to approximately solve any planning task defined over any environment in the original set, relying on the above interactions only. Can we design a provably efficient algorithm that achieves this ambitious goal of systematic generalization? In this paper, we give a partially positive answer to this question. First, we provide the first tractable formulation of systematic generalization by employing a causal viewpoint. Then, under specific structural assumptions, we provide a simple learning algorithm that allows us to guarantee any desired planning error up to an unavoidable sub-optimality term, while showcasing a polynomial sample complexity.

In the sequential decision making setting, an agent aims to achieve systematic generalization over a large, possibly infinite, set of environments. Such environments are modeled as discrete Markov decision processes with both states and actions represented through a feature vector. The underlying structure of the environments allows the transition dynamics to be factored into two components: one that is environment-specific and another one that is shared. Consider a set of environments that share the laws of motion as an illustrative example. In this setting, the agent can take a finite amount of reward-free interactions from a subset of these environments. The agent then must be able to approximately solve any planning task defined over any environment in the original set, relying on the above interactions only. Can we design a provably efficient algorithm that achieves this ambitious goal of systematic generalization? In this paper, we give a partially positive answer to this question. First, we provide the first tractable formulation of systematic generalization by employing a causal viewpoint. Then, under specific structural assumptions, we provide a simple learning algorithm that allows us to guarantee any desired planning error up to an unavoidable sub-optimality term, while showcasing a polynomial sample complexity.

Introduction
Whereas recent breakthroughs have established Reinforcement Learning (RL) (Sutton & Barto, 2018) as a powerful tool to address a wide range of sequential decision making problems, the curse of generalization (Kirk et al., 2021 still a main limitation of commonly used techniques. RL algorithms deployed on a given task are usually effective in discovering the correlation between an agent's behavior and the resulting performance from large amounts of labeled samples (Jaksch et al., 2010;Lange et al., 2012). However, those algorithms are usually unable to discover basic cause-effect relations between the agent's behavior and the environment dynamics. Crucially, the aforementioned correlations are oftentimes specific to the task at hand, and they are unlikely to be of any use for addressing different tasks or environments. Instead, some universal causal relations generalize over the environments, and once learned can be exploited for solving any task. Let us consider as an illustrative example an agent interacting with a large set of physical environments. While each of these environments can have its specific dynamics, we expect the basic laws of motion to hold across the environments, as they encode general causal relations. Once they are learned, there is no need to discover them again from scratch when facing a new task, or an unseen environment. Even if the dynamics over these relations can change, such as moving underwater is different than moving in the air, or the gravity can change from planet to planet, the underlying causal structure still holds. This knowledge alone often allows the agent to solve new tasks in unseen environments by taking a few, or even zero, interactions.
We argue that we should pursue this kind of generalization in RL, which we call systematic generalization, where learning universal causal relations from interactions with a few environments allows us to approximately solve any task in any other environment without further interactions. Although this problem setting might seem overly ambitious or even far-fetched, in this paper we provide the first tractable formulation of systematic generalization, thanks to a set of structural assumptions that are motivated by a causal viewpoint. The problem formulation is partially inspired by reward-free RL (Jin et al., 2020), in which the agent can take unlabelled interactions with an environment to learn a model that allows approximate planning for any reward function. Here, we extend this formulation to a large, potentially infinite, set of reward-free environments, or a universe, the agent can freely interact with. We consider discrete environments, such that both their states and actions can be described through vectors of discrete features. Crucially, these environments share a common causal structure that explains a significant portion, but not all, of their transition dynamics. Can we design a provably efficient algorithm that guarantees an arbitrarily small planning error for any possible task that can be defined over the set of environments, by taking reward-free interactions with a generative model?
In this paper, we provide a partially positive answer to this question by presenting a simple but principled causal modelbased approach (see Fig. 1). This algorithm interacts with a finite subset of the universe to learn the causal structure underlying the set of environments in the form of a causal dependency graph G (Wadhwa & Dong, 2021). Then, the causal transition model, which encodes the dynamics that is common across the environment, is obtained by estimating the Bayesian network P G over G (Dasgupta, 1997) from a mixture of the environments. Finally, the causal transition model is employed by a planning oracle to provide an approximately optimal policy for a latent environment and a given reward function. We can show that this simple recipe, with a sample complexity that is polynomial in all the relevant quantities of the problem, allows achieving any desired planning error up to an unavoidable error term. The latter is inherent to the setting, which demands generalization over an infinite set of environments, and cannot be overcome without additional samples from the test environment.
The sample complexity rates that we derived are detailed in Table 1, while the main contributions of this paper include: • The first tractable formulation of the systematic generalization problem in RL, thanks to structural assumptions motivated by causal considerations (Section 3); • A provably efficient algorithm to learn systematic generalization over an infinite set of environments with polynomial sample complexity (Section 4.1); • The sample complexity of estimating the causal structure underlying a discrete MDP (Section 4.2); • The sample complexity of estimating the Bayesian network underlying a discrete MDP (Section 4.3).

Preliminaries
We start with some basic notions about graphs, causality, and Markov decision processes that will be useful thereafter. We will denote a set of integers We will denote by U A the uniform distribution over A.

Graphs
We define a graph G as a pair G := (V, E), where V is a set of nodes and E ⊆ N × N is a set of edges between them. We call G a directed graph if all of its edges E are directed (i.e., ordered pairs of nodes). We also define the in-degree of a node to be its number of incoming edges: degree in (A) = |{(B, A) : (B, A) ∈ E, ∀B}|. G is said to be a Directed Acyclic Graph (DAG) if it is a directed graph without cycles. We call G a bipartite graph if there exists a partition X ∪ Y = V such that none of the nodes in X and Y are connected by an edge, i.e., E ∩ ( The skeleton of a graph G is the undirected graph that is obtained from G by replacing all the directed edges in E with undirected ones. Finally, the graph edit distance between two graphs is the minimum number of graph edits (addition or deletion of either a node or an edge) necessary to transform one graph into the other.

Causal Graphs and Bayesian Networks
For a set X of random variables, we represent the causal structure over X with a DAG G X = (X , E), 1 which we call the causal graph of X . For each pair of variables A, B ∈ X , a directed edge (A, B) ∈ G X denotes that B is conditionally dependent on A. For every variable A ∈ X , we denote as Pa(A) the causal parents of A, i.e., the set of all the variables B ∈ X on which A is conditionally dependent, (B, A) ∈ G X .
A Bayesian network (Dean & Kanazawa, 1989) over the set X is defined as N := (G X , P ), where G X specifies the structure of the network, i.e., the dependencies between the variables in X , and the distribution P : X → ∆ X specifies the conditional probabilities of the variables in X , such that

Markov Decision Processes
A tabular episodic Markov Decision Process (MDP) (Puterman, 2014) is defined as M := (S, A, P, H, r), where S is a set of |S| = S states, A is a set of |A| = A actions, P is a transition model such that P (s |s, a) gives the conditional probability distribution of the next state s having taken action a in state s, H is the horizon of an episode, and r : S × A → [0, 1] is a deterministic reward function.
The strategy of an agent interacting with M is represented by a non-stationary, stochastic policy, a collection of functions (π h : S → ∆ A ) h∈ [H] where π h (a|s) denotes the conditional probability of taking action a in state s at step h. The value function V π h : S → R associated to π is defined as the expected sum of the rewards that will be collected, under the policy π, starting from s at step h, i.e., For later convenience, we further define P V π h+1 (s, a) := E s ∼P (·|s,a) [V π h+1 (s )] and V π 1 := E s∼P [V π 1 (s)]. We will write V π M,r to denote V π 1 in the MDP M with reward function r (if not obvious from the context). For an MDP M with finite states, actions, and horizon, there always exists an optimal policy π * that gives the value V * h (s) = sup π V π h (s) for every s, a, h. The goal of the agent is to find a policy π that is -close to the optimal one, i.e., V * 1 − V π 1 ≤ . Finally, we define a discrete Markov decision process as M := ((S, d S , n), (A, d A , n), P, H, r), where S, A, P, H, r are specified as before, and where the states and actions spaces admit additional structure, such that every s ∈ S can be represented through a d S -dimensional vector of discrete features taking value in [n], and every a ∈ A can be represented through a d A -dimensional vector of discrete features taking value in [n]. Note that any tabular MDP can be formulated under this alternative formalism by taking n = 2, d S = S, and d A = A, so that the states and actions are specified through one-hot encoding.

Problem Formulation
In our setting, a learning agent aims to master a large, potentially infinite, set U of environments modeled as discrete MDPs without rewards, , which we call a universe. The agent can draw a finite amount of experience by interacting with the MDPs in U. From these interactions alone, the agent aims to acquire sufficient knowledge to approximately solve any task that can be specified over the universe U. Specifically, a task is defined as any pairing of an MDP M ∈ U and a reward function r, whereas solving it refers to providing a slightly sub-optimal policy via planning, i.e., without taking additional interactions. We call this problem systematic generalization, which we can formalize as follows. Definition 1 (Systematic Generalization). For any latent MDP M ∈ U and any given reward function r : S × A → [0, 1], the systematic generalization problem requires the agent to provide a policy π, such that V * M,r − V π M,r ≤ up to any desired sub-optimality > 0.
Since the set U is infinite, we clearly require additional structure to make the problem feasible. On the one hand, the state space (S, d S , n), action space (A, d A , n), and initial state distribution µ are shared across M ∈ U. The transition dynamics P i is instead specific to each MDP M i ∈ U. However, we assume the presence of a common causal structure that underlies the transition dynamics of the universe, and relates the single transition models P i .

The Causal Structure of the Transition Dynamics
The transition dynamics of a discrete MDP gives the conditional probability of next state features s given the current [n] d S + d A Figure 2. Illustration of the causal structure G of U.
state-action features (s, a). To ease the notation, from now on we will denote the current state-action features with a random vector X = ( , and the next state features with a random For each environment M i ∈ U, the conditional dependencies between the next state features Y and the current state-action features X are represented through a bipartite dependency graph . Clearly, each environment can display its own dependencies, but we assume there is a set of dependencies that represent general causal relationships between the features, and that appear in any M i ∈ U. In particular, we call the intersection G := ∩ ∞ i=0 G i the causal structure of U, which is the set of conditional dependencies that are common across the universe. In Figure 2, we show an illustration of such a causal structure. Since it represents universal causal relationships, the causal structure G is time-consistent, i.e., G (h) = G (1) for any step h ∈ [H], and we further assume that G is sparse, which means that the number of features X[z] on which a feature Y [j] is dependent on is bounded from above.
Given a causal structure G, and without losing generality, 2 we can express each transition model P i as in which P G is the Bayesian network over the causal structure G, whereas F i includes environment-specific factors affecting the conditional probabilities. 3 Since it represents the conditional probabilities due to universal causal relations in U, we call P G the causal transition model of U. Thanks to the structure G, this model can be further factored as where the scopes Z j are the set of indices z such that . For the purpose of this work, we make the crucial assumption that the causal transition model is non-vacuous and that it explains a significant part of the transition dynamics of M i ∈ U. Assumption 2 (λ-sufficiency). Let λ ∈ [0, 1] be a constant. The causal transition model P G is causally λ-sufficient if Notably, the parameter λ controls the amount of the transition dynamics that is due to the universal causal relations G (λ = 0 means that P G is sufficient to explain the transition dynamics of any M i ∈ U, whereas λ = 1 implies no shared structure between the transition dynamics of the M i ∈ U).
In this paper, we argue that learning the causal transition model P G is a good target for systematic generalization and we provide theoretical support for this claim in Section 4.

A Class of Training Environments
Even if the universe U admits the structure that we presented in the last section, it is still an infinite set. Instead, the agent can only interact with a finite subset of discrete MDPs which we call a class of size M . Crucially, the causal structure G is a property of the full set U, and if we aim to infer it from interactions with a finite class M, we have to assume that M is informative enough on the universal causal relations of U.
Analogously, if we aim to infer the causal transition model P G from interactions with the transition models P i of the single MDPs M i ∈ M, we have to assume that M is balanced in terms of the conditional probabilities displayed by its components, so that the factors that do not represent universal causal relations even out while learning.
Assumption 4 (Evenness). Let M ⊂ U be class of size M . We say that M is causally even if Whereas in this paper we assume that M is diverse and even by design, we leave as future work the interesting problem of selecting such a class from active interactions with U, which would add to our problem formulation flavors of active learning and experimental design (Hauser & Bühlmann, 2014;Kocaoglu et al., 2017;Ghassami et al., 2018).

Learning Systematic Generalization
Before addressing the sample complexity of systematic generalization, it is worth considering the kind of interactions that we need in order to learn the causal transition model P G and its underlying causal structure G. Especially, thanks to the peculiar configuration of the causal structure G, i.e., a bipartite graph in which the edges are necessarily directed from the state-action features X to the next state features Y , as a causation can only happen from the past to the future, learning the skeleton of G is equivalent to learning its full structure. Crucially, learning the skeleton of a causal graph does not need specific interventions, as it can be done from observational data alone (Hauser & Bühlmann, 2014).
Proposition 1. The causal structure G of U can be identified from purely observational data.
In this paper, we will consider the online learning setting with a generative model for estimating G and P G from sampled interactions with a class M of size M . A generative model allows the agent to set the state of an MDP before sampling a transition, instead of drawing sequential interactions from the process. Finally, analogous results to what we obtain here can apply to the offline setting as well, in addition to convenient coverage assumptions on the dataset.

Sample Complexity Analysis
We provide a sample complexity analysis of the problem, which stands as a core contribution of this paper along with the problem formulation itself that we provided in Section 3.
First, we consider the sample complexity of systematic generalization (Section 4.1). Then, we provide ancillary sample complexity results on the estimation of the causal structure (Section 4.2) and the Bayesian network (Section 4.3) of an MDP, which can be of independent interest.

Sample Complexity of Systematic Generalization with a Generative Model
We have access to a class M of discrete MDPs within a universe U, from which we can draw interactions with a generative model P (X). We would like to solve the systematic generalization problem as described in Definition 1. This problem requires to provide, for any combination of a (latent) MDP M ∈ U, and a given reward function r, a planning policy π such that V * M,r − V π M,r ≤ . Especially, can we design an algorithm that guarantees this requirement with high probability by taking a number of samples K that is polynomial in and the relevant parameters of M?
Here we give a partially positive answer to this question, by providing a simple but provably efficient algorithm that guarantees systematic generalization over U up to an unavoidable sub-optimality term λ that we will later specify.
The algorithm implements a model-based approach into two separated components. The first, for which we provide the pseudocode in Algorithm 1, is the procedure that actually interacts with the class M to obtain a principled estimation P G of the causal transition model P G of U. The second, is a planning oracle that takes as input a reward function r and the estimated causal transition model, and returns an optimal policy π operating on P G as an approximation of the transition model P i of the true MDP M i . 5 First, we provide the sample complexity of the causal transition model estimation (Algorithm 1), which in turn is based on repeated causal structure estimations (Algorithm 2) to obtain G, and an estimation procedure of the Bayesian network over G (Algorithm 3) to obtain P G .
be a class of M discrete MDPs, let δ ∈ (0, 1), and let > 0. The Algorithm 1 returns a causal transition model P G such that P r( P G − P G 1 ≥ ) ≤ δ with a sample complexity An analogous result can be derived for the tabular MDP setting, as stated by the following lemma.
MDPs. The sample complexity of Lemma 4.1 reduces to Having established the sample complexity of the causal transition model estimation, we can now show how the learned model P G allows us to approximately solve, via a planning oracle, any task defined by a combination of a latent MDP M i ∈ U and a given reward function r. To provide this result in the discrete MDP setting, we have to further assume that the transition dynamics P i of the target MDP M i admits factorization analogous to (1), such that we can write where the scopes Z j are given by the environment-specific causal structure G i , which we assume to be 2Z-sparse (Assumption 1).
Theorem 4.3. Let δ ∈ (0, 1) and > 0. For a latent discrete MDP M ∈ U, and a given reward function r, a planning oracle operating on the causal transition model P G as an approximation of M returns a policy π such that Without the additional factorization of the environmentspecific transition model, the result of Theorem 4.3 reduces to the analogous for the tabular MDP setting. Theorem 4.3 and Corollary 4.4 establish the sample complexity of systematic generalization through Lemma 4.1 and Lemma 4.2 respectively. For the discrete MDP setting, we have that O(M H 6 d 3 S Z 2 n 5Z+3 ), which reduces to O(M H 6 S 4 A 2 Z 2 ) in the tabular setting. Unfortunately, we are only able to obtain systematic generalization up to an unavoidable sub-optimality term λ . This error term is related to the λ-sufficiency of the causal transition model (Assumption 2), and it accounts for the fact that P G cannot fully explain the transition dynamics of each M ∈ U, even when it is estimated exactly. This is inherent to the ambitious problem setting, and can be only overcome with additional interactions with the test MDP M.

Sample Complexity of Learning the Causal Structure of a Discrete MDP
As a byproduct of the main result in Theorem 4.3, we can provide a specific sample complexity result for the problem of learning the causal structure G underlying a discrete MDP M with a generative model. We believe that this problem can be of independent interest, mainly in consideration of

Algorithm 1 Causal Transition Model Estimation
Input: class of MDPs M, error , confidence δ Output: causal transition model P G previous work on causal discovery of general stochastic processes (e.g., Wadhwa & Dong, 2021), for which we refine known results to account for the structure of an MDP, which allows for a tighter analysis of the sample complexity.
Instead of targeting the exact dependency graph G, which can include dependencies that are so weak to be nearly impossible to detect with a finite number of samples, we only address the dependencies above a given threshold.
Before presenting the sample complexity result, we state the existence of a principled independence testing procedure.
We can now provide an upper bound to the number of samples required by a simple estimation procedure to return an ( , δ)-estimate G of the causal dependency graph G.
The Algorithm 2 returns a dependency graph G such that P r( G = G ) ≤ δ with a sample complexity Finally, we can state an analogous result for a tabular MDP setting, by taking n = 2, d S = S, d A = A. Corollary 4.7. Let M be a tabular MDP. The result of Theorem 4.6 reduces to K = O log(S 2 A/δ)/ 2 .

Sample Complexity of Learning the Bayesian Network of a Discrete MDP
We present as a standalone result an upper bound to the sample complexity of learning the parameters of a Bayesian network P G with a fixed structure G. Especially, we refine known results (e.g., Friedman & Yakhini, 1996;Dasgupta, 1997;Cheng et al., 2002;Abbeel et al., 2006;Canonne et al., 2017) by considering the specific structure G of an MDP.
If the structure G is dense, the number of parameters of P G grows exponentially, making the estimation problem mostly intractable. Thus, we consider a Z-sparse G (Assumption 1), as in previous works (Dasgupta, 1997). Then, we can provide a polynomial sample complexity for the problem of learning the Bayesian network P G of a discrete MDP M.
The Algorithm 3 returns a Bayesian network P G such that P r( P G − P G 1 ≥ ) ≤ δ with a sample complexity Analogously, we can state the sample complexity of learning the fixed structure Bayesian network of a tabular MDP.
Corollary 4.9. Let M be a tabular MDP. The result of Theorem 4.8 reduces to

Numerical Validation
We empirically validate the theoretical findings of this work by experimenting on a synthetic example where each environment is a person, and the MDP represents how a series of actions the person can take influences their weight (W ) and academic performance (A). As actions we consider hours of physical training (P ), hours of sleep (S), hours of study (St), amount of vegetables in the diet (D), and the amount of caffeine intake (C). The obvious use-case for such a model would be a tracking device that monitors how the actions of a person influence their weight and academic performance and provides personalized recommendations to

Algorithm 3 Bayesian Network Estimation for an MDP
Output: Bayesian network P G reach the person's goals. While the physiological responses of different individuals can vary, there are some underlying mechanisms shared by all humans, and therefore deemed causal in our terminology. Examples of such causal links are the dependency of weight on the type of diet, and the dependency of academic performance on the number of hours of study. Other links, such as the dependency of weight on the amount of caffeine can be present in some individuals, but are generally not shared and therefore not causal. We report in the Appendix ( Figure 5) an illustration of the described domain. For simplicity, all variables are treated as discrete with possible values 0 (below average), 1 (average) or 2 (above average). See Appendix B for details on how transition models of different environments are generated. A class M of 3 environments is used to estimate the causal transition model. All experiments are repeated 10 times and report the mean and standard deviation of the desired quantities.
Causal Structure Estimation We first empirically investigate the graph edit distance between estimated and groundtruth causal structure, GED(G, G), as a function of number of samples (K in Algorithm 1). The causal structure is estimated by obtaining the causal graph for each training environment (using a series of independence tests), and taking the intersection of their edges. As expected, as we increase the number of samples, the distance converges to zero, and we can perfectly recover the causal graph (Figure 4a). Figure 4b presents the L 1 -distance between the estimate and ground truth causal transition model, as a function of the total number of samples (K + K in Algorithm 1). As the number of samples grows, the L 1 -distance shrinks, converging at around 0.05, an estimation error due to the environments  not respecting the evenness assumption.

Causal Transition Model Estimation
Value Function Estimation Finally, we investigate how well we can approximate the optimal value function for a new unseen environment (Figure 4c). We observe that our algorithm is able to effectively approximate the optimal value function up to a small error with a reasonable number of samples.

Related Work
We revise the relevant literature and we discuss how it relates with our problem formulation and the reported results. whenever U is small, but leveraging the inner structure of U becomes more and more important as U grows to infinity, while M remains constant. However, our approach pays this further generality with the additional error term λ , which is unavoidable. Finally, it is an interesting direction for future work to establish whether the additional factors in S, A, H w.r.t. reward-free RL are also unavoidable.

Reward-Free Reinforcement Learning
Multi-Task Reinforcement Learning In the multi-task RL setting (Lazaric & Ghavamzadeh, 2010;Lazaric & Restelli, 2011;Brunskill & Li, 2013;Liu et al., 2016), the agent interacts with a set of MDPs (tasks) pursuing generalization from one task to the other, which is in a way similar to the aim of systematic generalization. Although sample complexity results have been derived for the multitask setting (e.g., Brunskill & Li, 2013), we are not aware of previous works in this stream that considered the problem of interacting with a set of reward-free MDPs, nor the presence of a common causal structure underlying the set of MDPs.

Model-Based Reinforcement Learning
The modelbased RL (Sutton & Barto, 2018) methodology prescribes learning an approximate model of transition dynamics in order to learn an optimal policy. Theoretical works (e.g., Jaksch et al., 2010;Ayoub et al., 2020) generally concern with the estimation of the approximate value functions obtained through the learned model, rather than the estimation of the model itself. A notable exception is the work by Tarbouriech et al. (2020), which targets point-wise high probability guarantees on the model estimation as we do in Lemma 4.1, 4.2. However, they address the model estimation of a single tabular MDP M, instead of the shared transition dynamics of an infinite set of MDPs U that we target in this paper. Other works (e.g., Zhang et al., 2020;Tomar et al., 2021) have also addressed model-based RL from a causal perspective. To the best of our knowledge, we are the first to prove a polynomial sample complexity for causal model-based RL in systematic generalization.

Factored Markov Decision Processes
The factored MDP formalism (Kearns & Koller, 1999) allows encoding transition dynamics that are the product of multiple independent factors. This formalism is closely related to how we define the causal transition model as a product of independent factors in (1), which can be actually seen as a factored MDP.

Conclusion
In this work, we formulate the problem of systematic generalization in reinforcement learning from a causal viewpoint. We then propose a provably efficient algorithm which partially solves the problem under certain structural assumptions motivated by causal considerations.

A. Proofs
Proofs of Section 3.3 Proposition 1. The causal structure G of U can be identified from purely observational data.
Proof. First, recall that with observational data alone, a causal graph can be identified up to its Markov equivalence class (Hauser & Bühlmann, 2014). This means that its skeleton and v-structure are properly identified, meanwhile determining the edge orientations requires interventional data in the general case. Since in the considered causal graph G the edges orientations are determined a priori (as they follow the direction of time), the causal graph can be entirely determined by using only observational data.

Proofs of Section 4.1: Causal Transition Model Estimation
Before reporting the proof of the main result in Theorem 4.3, it is worth considering a set of lemmas that will be instrumental to the main proof.
First, we provide an upper bound to the L1-norm between the Bayesian network P G over a given structure G and the Bayesian network P G over the structure G , which is the -dependency subgraph of G as defined in Definition 2.
Lemma A.1. Let G a Z-sparse dependency graph, and let G its corresponding -dependence subgraph for a threshold > 0. The L1-norm between the Bayesian network P G over G and the Bayesian network P G over G can be upper bounded as Proof. The proof is based on the fact that every edge (X i , Y j ) such that (X i , Y j ) ∈ G and (X i , Y j ) / ∈ G corresponds to a weak conditional dependence (see Definition 2), which means that P Yj |Xi − P Yj 1 ≤ .
We denote with Z j the scopes of the parents of the node Y [j] in G, i.e., Pa G (Y [j]) = X[Z j ], and with Z j, the scopes of the parents of the node Y [j] in G , i.e., Pa G (Y [j]) = X[Z j, ]. As a direct consequence of Definition 2, we have Z j, ⊆ Z j for any j ∈ d S , and we can write Then, we let Z j \ Z j, = [I] overwriting the actual indices for the sake of clarity, and we derive in which we employed the property µ − ν 1 ≤ i µ i − i ν i 1 ≤ i µ i − ν i 1 for the L1-norm between product distributions µ = i µ i , ν = i ν i to write (2), we repeatedly applied the triangle inequality µ−ν 1 ≤ µ−ρ 1 + ρ−ν 1 to get (3) from (2), we upper bounded each term of the sum in (3) with thanks to Definition 2, and we finally employed the Z-sparseness Assumption 1 to upper bound I with Z in (4).
Next, we provide a crucial sample complexity result for a provably efficient estimation of a Bayesian network P G over an estimated -dependency subgraph G, which relies on both the causal structure estimation result of Theorem 4.6 and the Bayesian network estimation result of Theorem 4.8.
Lemma A.2. Let M be a discrete MDP, let M = {M} be a singleton class, let δ ∈ (0, 1), and let > 0. The Algorithm 1 returns a Bayesian network P G such that P r( P G − P G 1 ≥ ) ≤ δ with a sample complexity Proof. We aim to obtain the number of samples K = K + K for which Algorithm 1 is guaranteed to return a Bayesian network estimate P G over a causal structure estimate G such that P r( P G − P G 1 ≥ ) ≤ δ in a setting with a singleton class of discrete MDPs. First, we derive the following decomposition of the error in which we employed the triangle inequality µ − ν 1 ≤ µ − ρ 1 + ρ − ν 1 . Then, we can write Bayesian network estimation ( ) Bayesian network subgraph ( ) through the decomposition (5) and a union bound to isolate the three independent sources of error ( ), (•), ( ). To upper bound the latter term ( ) with 0, we invoke Lemma A.1 to have d s Z ≤ 3 , which gives ≤ 3d S Z . Then, we consider the middle term (•), for which we can write We can now upper bound (•) ≤ δ/2 through (6) by invoking Theorem 4.6 with threshold = 3d S Z and confidence δ = δ 2 , which gives Next, we can upper bound ( ) ≤ δ/2 by invoking Theorem 4.8 with threshold = 3 and confidence δ = δ 2 , which gives Finally, through the combination of (7) and (8), we can derive the sample complexity that guarantees P r( P G − P G 1 ≥ ) ≤ δ under the assumption 4/3 2 , i.e., which concludes the proof.
Whereas Lemma A.2 is concerned with the sample complexity of learning the Bayesian network of a singleton class, we can now extend the result to account for a class M composed of M discrete MDPs.
be a class of M discrete MDPs, let δ ∈ (0, 1), and let > 0. The Algorithm 1 returns a causal transition model P G such that P r( P G − P G 1 ≥ ) ≤ δ with a sample complexity Proof. We aim to obtain the number of samples K = M K + K for which Algorithm 1 is guaranteed to return a Bayesian network estimate P G over a causal structure estimate G such that P r( P G − P G 1 ≥ ) ≤ δ in a setting with a class of M discrete MDPs. First, we can derive an analogous decomposition as in (5), such that we have P r P G − P G 1 ≥ ≤ P r P G − P G 1 ≥ 3 Bayesian network estimation ( ) + P r P G − P G 1 ≥ 3 causal structure estimation (•) + P r P G − P G 1 ≥ 3 Bayesian network subgraph ( ) through a union bound. Crucially, the terms ( ), ( ) are unaffected by the class size, which leads to K = (8) by upper bounding ( ), and ≤ 3d S Z by upper bounding ( ), exactly as in the proof of Lemma A.2. Instead, the number of samples K has to guarantee that (•) = P r( P G − P G 1 ≥ /3) ≤ δ/2, where the causal structure G is now the intersection of the causal structures of the single class components M i , i.e., G = ∩ M i=1 G ,i . Especially, we can write through a union bound on the estimation of the single causal structures G i . Then, we can upper bound (•) ≤ δ/2 through (9) by invoking Theorem 4.6 with threshold = 3d S Z and confidence δ = δ 2M , which gives Finally, through the combination of (10) and (8), we can derive the sample complexity that guarantees P r( P G − P G 1 ≥ ) ≤ δ under the assumption 4/3 2 , i.e., which concludes the proof.
It is now straightforward to extend Lemma 4.1 for a class M composed of M tabular MDPs.
Proof. To obtain K = M K + K , we follows similar steps as in the proof of Lemma 4.1, to have the usual decomposition of the event P r( P G − P G 1 ≥ ) in the ( ), (•), ( ) terms. We can deal with ( ) as in Lemma 4.1 to get ≤ 3SZ . Then, we upper bound (•) ≤ δ/2 by invoking Corollary 4.7 (instead of Theorem 4.6) with threshold = 3SZ and confidence δ = δ 2M , which gives Similarly, we upper bound ( ) ≤ δ/2 by invoking Corollary 4.9 (instead of Theorem 4.5) with threshold = 3 and confidence δ = δ 2 , which gives Finally, we combine 11 with 12 to obtain

Proofs of Section 4.1: Planning
For the sake of notational clarity within the following proofs we express Theorem 4.3. Let δ ∈ (0, 1) and > 0. For a latent discrete MDP M ∈ U, and a given reward function r, a planning oracle operating on the causal transition model P G as an approximation of M returns a policy π such that where λ = 2λH 3 d S n 2Z+1 , and P G is obtained from Algorithm 1 with δ = δ and = /2H 3 n Z+1 .
Proof. Consider the MDPs with transition model P and P G . We refer to the respective optimal policies as π * and π * . Moreover, since the reward r is fixed, we remove it from the expressions for the sake of clarity, and refer with V to the value function of the MDP with transition model P G . As done in (Jin et al., 2020, Theorem 3.5), we can write the following decomposition, where V * := V π * .
where in the last step we have set to 0 the approximation due to the planning oracle assumption, and we have bounded the evaluation errors according to Lemma A.3. In order to get 2n Z+1 H 3 = we have to set = 2n Z+1 H 3 . Considering the sample complexity result in Lemma 4.1 the final sample complexity will be: Under the preconditions of Theorem 4.3, with probability 1 − δ, for any reward function r and policy π, we can bound the value function estimation error as follows.
where V is the value function of the MDP with transition model P G , is the approximation error between P G and P G studied in Lemma 4.1, and λ stands for the λ-sufficiency parameter of P G .
Proof. The proof will be along the lines of that of Lemma 3.6 in (Jin et al., 2020). We first recall (Dann et al., 2017, Lemma E.15), which we restate in Lemma A.5. In this prof we consider an environment specific true MDP M with transition model P , and an mdp M that has as transition model the estimated causal transition model P G . In the following, the expectations will be w.r.t. P . Moreover, since the reward r is fixed, we remove it from the expressions for the sake of clarity. We can start deriving We now bound a single term within the sum above as follows: We can now bound each term. Let us start considering the first term: We can now bound the first term of (16): Due to the uniform sampling and Z-sparseness assumptions, we have P G (X[Z j ]) = 1 n Z , hence: = n Z Therefore: Replacing this in (17) and marginalizing over Y \Y [j] we obtain: Where d S is the approximation term of each component. By plugging this bound into (16) we get: where in the last step we have recursively bounded the right terms as in (24). By considering 2Z-sparseness, λ-sufficiency, and that the transition model P factorizes, we can apply the same procedure to bound the second term of equation (15) as: Therefore the initial expression in (14) becomes: Corollary 4.4. For a tabular MDP M ∈ M, the result of Theorem 4.3 holds with λ = 2λSAH 3 , = /2SAH 3 .
Proof. Consider the MDPs with transition model P and P G . We refer to the respective optimal policies as π * and π * . Moreover, since the reward r is fixed, we remove it from the expressions for the sake of clarity, and refer with V to the value function of the MDP with transition model P G . As done in (Jin et al., 2020, Theorem 3.5), we can write the following decomposition, where V * := V π * .
where in the last step we have set to 0 the approximation due to the planning oracle assumption, and we have bounded the evaluation errors according to Lemma A.4. In order to get 2SAH 3 = we have to set = 2SAH 3 . Considering the sample complexity result in Lemma 4.2 the final sample complexity will be: Lemma A.4. Under the preconditions of Corollary 4.4, with probability 1 − δ, for any reward function r and policy π, we can bound the value function estimation error as follows.
where V is the value function of the MDP with transition model P G , is the approximation error between P G and P G studied in Lemma 4.1, and λ stands for the λ-sufficiency parameter of P G .
Proof. The proof will be along the lines of that of Lemma 3.6 in (Jin et al., 2020). We first recall (Dann et al., 2017, Lemma E.15), which we restate in Lemma A.5. In this prof we consider an environment specific true MDP M with transition model P , and an mdp M that has as transition model the estimated causal transition model P G . In the following, the expectations will be w.r.t. P . Moreover, since the reward r is fixed, we remove it from the expressions for the sake of clarity. We can start deriving We now bound a single term within the sum above as follows: where in the last step we have used the fact that there must exist an optimal deterministic policy. Due to the uniform sampling assumption, we have P (s, a) = 1 SA , hence: max π † P π † (s, a) P (s, a) ≤ 1 P (s, a) = SA Therefore: P π † (s, a) ≤ SA · P (s, a) Moreover, notice that, since π is deterministic we have P π (s) = P π (s) = P π (s, a) ≤ SA · P (s, a). Replacing it in the expression above we get where is the approximation error between P G and P G studied in Lemma 4.1, and in the penultimate step we have used the following derivation: Hence, due to this recursive unrolling, we have: Notice that the same argument holds also for the second term of (22), replacing with λ.
By plugging the result in equation (23) into the initial expression we get: In the following we restate (Dann et al., 2017, Lemma E.15) for the case of stationary transition model.
Lemma A.5. For any two MDPs M and M with rewards r and r and transition models P and P , the difference in value functions V , V w.r.t. the same policy π can be written as: Proofs of Section 4.2 We provide the proof of the sample complexity result for learning the causal structure of a discrete MDP with a generative model.
Theorem 4.6. Let M be a discrete MDP with an underlying causal structure G, let δ ∈ (0, 1), and let > 0. The Algorithm 2 returns a dependency graph G such that P r( G = G ) ≤ δ with a sample complexity Proof. We aim to obtain the number of samples K for which Algorithm 2 is guaranteed to return a causal structure estimate G such that P r( G = G ) ≤ δ in a discrete MDP setting. First, we can upper bound the probability of the bad event P r( G = G ) in terms of the probability of a failure in the independence testing procedure I(X z , Y j ) for a single pair of nodes X z ∈ G , Y z ∈ G , i.e., where we applied an union bound to obtain the last inequality. Now we can look at the probability of a single independence test failure. Especially, for a provably efficient independence test (the existence of such a test is stated by Lemma 4.5, whereas the Algorithm 2 in Diakonikolas et al. (2021) reports an actual testing procedure), we have P r(test I(X z , Y j ) fails) ≤ δ , for any choice of δ ∈ (0, 1), > 0, with a number of samples  (27), we obtain P r( G = G ) with a sample complexity under the assumption 2 4/3 , which concludes the proof.
The proof of the analogous sample complexity result for a tabular MDP setting (Corollary 4.7) is a direct consequence of Theorem 4.6 by letting n = 2, d S = S, d A = A.

Proofs of Section 4.3
We first report a useful concentration inequality for the L1-norm between the empirical distribution computed over K samples and the true distribution (Weissman et al., 2003, Theorem 2.1). Lemma A.6 (Weissman et al. (2003)). Let X 1 , . . . , X K be i.i.d. random variables over [n] having probabilities P r(X k = i) = P i , and let P K (i) = 1 K K k=1 1(X k = i). Then, for every threshold > 0, it holds P r P K − P 1 ≥ ≤ 2 exp(−K 2 /2n).
We can now provide the proof of the sample complexity result for learning the Bayesian network of a discrete MDP with a given causal structure. Theorem 4.8. Let M be a discrete MDP, let G be its underlying causal structure, let δ ∈ (0, 1), and let > 0. The Algorithm 3 returns a Bayesian network P G such that P r( P G − P G 1 ≥ ) ≤ δ with a sample complexity Proof. We aim to obtain the number of samples K for which Algorithm 3 is guaranteed to return a Bayesian network estimate P G such that P r( P G − P G 1 ≥ ) ≤ δ in a discrete MDP setting. First, we note that in which we employed the property µ − ν 1 ≤ i µ i − i ν i 1 ≤ i µ i − ν i 1 for the L1-norm between product distributions µ = i µ i , ν = i ν i to write (28), and we applied a union bound to derive (31) from (30). Similarly, we can write P r P j − P j 1 ≥ d S ≤ P r x∈[n] |Z j | P j (·|x) − P j (·|x) 1 ≥ d S n |Zj | (32) ≤ x∈[n] |Z j | P r P j (·|x) − P j (·|x) 1 ≥ d S n |Zj | (33) by applying a union bound to derive (33) from (32), and by employing Assumption 1 to bound |Z j | with Z in (34). We can now invoke Lemma A.6 to obtain the sample complexity K that guarantees P r( P j (·|x) − P j (·|x) 1 ≥ ) ≤ δ , i.e., K = 2n log(2/δ ) ( ) 2 = 2 d 2 S n 2Z+1 log(2d S n Z /δ) 2 , where we let = d S n Z , δ = δ d S n Z . Finally, by summing K for any x ∈ [nm] |Zj | and any j ∈ [d S ], we obtain which proves the theorem.
To prove the analogous sample complexity result for a tabular MDP we can exploit a slightly tighter concentration on the KL divergence between the empirical distribution and the true distribution in the case of binary variables (Dembo & Zeitouni, 2009, Theorem 2.2.3) 6 , which we report for convenience in the following lemma.
We can now provide the proof of Corollary 4.9.
Proof. We aim to obtain the number of samples K for which Algorithm 3 is guaranteed to return a Bayesian network estimate P G such that P r( P G − P G 1 ≥ ) ≤ δ in a tabular MDP setting. We start by considering the KL divergence d KL ( P G ||P G ). Especially, we note Then, for any > 0 we can write P r d KL P G ||P G ≥ ≤ P r in which we applied a first union bound to get (36) from (35), a second union bound to get (38) from (37), and Assumption 1 to bound |Z j | with Z in (39). We can now invoke Lemma A.7 to obtain the sample complexity K that guarantees P r(d KL ( P j (·|x)||P j (·|x)) ≥ ) ≤ δ , i.e., where we let = S2 Z , and δ = δ S2 Z for any choice of δ ∈ (0, 1). By summing K for any x ∈ [2] |Zj | and and j ∈ [S], we obtain the sample complexity K that guarantees P r(d KL ( P G ||P G ) ≥ ) ≤ δ , i.e., Finally, we employ the Pinsker's inequality P G − P G 1 ≤ 2d KL ( P G ||P G ) (Csiszár, 1967) to write which gives the sample complexity K that guarantees P r( P G − P G 1 ≥ ) ≤ δ by letting = 2 2 and δ = δ in (40), i.e.,

B. Numerical Validation
Let A, W , P , S, D, C and St represent academic performance, weight, physical activity, sleep, diet and study respectively. We start by defining the causal transition model as follows: The transition model for a specific environment is then generated by adding independent white noise to each coefficient in the above equations for µ A and µ W , including the ones not shown because set to 0: µ A = (1 + 1 )A + 2 W + 3 P + (0.2 + 4 )D + (0.5 + 5 )S + 6 C + (0.8 + 7 )St − 0.8 µ W = 8 A + (1 + 9 )W + (−0.5 + 10 )P + (−0.5 + 11 )D + 12 S + 13 C + 14 St + 0. causal edge e ∈ G environment specific edge e / ∈ G Figure 5. Bayesian network for a specific environment (person) of our synthetic example. Grey edges are causal and therefore shared by other environments in the same universe U, while blue dashed edges are environment-specific dependencies.


Discriminative feature generation for classification of imbalanced data

The data imbalance problem is a frequent bottleneck in the classification performance of neural networks. In this paper, we propose a novel supervised discriminative feature generation (DFG) method for a minority class dataset. DFG is based on the modified structure of a generative adversarial network consisting of four independent networks: generator, discriminator, feature extractor, and classifier. To augment the selected discriminative features of the minority class data by adopting an attention mechanism, the generator for the class-imbalanced target task is trained, and the feature extractor and classifier are regularized using the pre-trained features from a large source data. The experimental results show that the DFG generator enhances the augmentation of the label-preserved and diverse features, and the classification results are significantly improved on the target task. The feature generation model can contribute greatly to the development of data augmentation methods through discriminative feature generation and supervised attention methods.

Introduction
Deep learning has achieved significant improvements in various tasks in computer vision applications using open image datasets containing a large amount of data. However, the acquisition of large datasets is a challenge in real-world applications, especially if they pertain to new areas in deep learning. For training deep neural networks, the quantity of data is as important as the quality because training with only a small dataset often results in degradation of the neural network performance or introduces an overfitting problem. Attempts have been made to establish a large training dataset in various fields; however, the annotation of a dataset still remains expensive, laborious, and timeconsuming. Furthermore, the distribution of classes in the dataset is often imbalanced. Training a network under class-imbalanced conditions can produce a detrimental effect on neural network performance and a biased classification result [1,2,3].
Data augmentation enhances the size and quality of the training dataset. Techniques range from simple data augmentation, such as flip, shift, and rotation, to deep generative models based on generative adversarial networks (GAN) [4], where a small training dataset is augmented at the input data level. Because GANs can approximate the distribution of the real input data and generate realistic samples from a generative model, many recent studies have shown that augmenting a small training dataset using a GAN can improve the classification performance in real applications [5,6,7,8,9,10]. However, such data augmentation at input level showed the limited improvement in the performance, because balancing the data distribution have week relation to the enhancement of feature extraction in the minority dataset. Also, the state-of-art GAN has the limited generating capability on large-scale images [11,12]. When the size of input image is large, these augmentation method is not applicable.
Unlike data augmentation, adversarial feature augmentation generates domain-invariant features, increasing the size of the minority classes in the feature space without considering the modality of the input data [13]. This augmentation was adopted in supervised learning, especially for the class imbalance problem [14]. Although the feature augmentation method performs better than the unsupervised domain adaptation method by generating domain-invariant and modality-free features, the generated feature maps from the outer fully connected layer do not have sufficient feature dimensions to represent the data distribution with a small amount of minority class data. In other words, generating features in convolutional layers has a greater capacity to improve classification performance in the case of the class-imbalanced conditions.
Taking these problems into consideration, we propose a novel discriminative feature generation (DFG) method using attention maps in the feature space. The proposed method is a combination of transfer learning and adversarial feature augmentation to complement their drawbacks. The baseline is transfer learning with regularizations, and the features are augmented using a GAN with the weight of the activation level of the feature maps in each class.
The framework is based on a modified GAN structure containing an independent classifier to improve the neural network performance. We extend the adversarial feature augmentation [15] in the following three manners: (1) Our GAN structure is composed of four independent networks: a feature generator, feature discriminator, feature extractor, and feature classifier. The more constituent networks enable the network to be trained for generating more effective features, especially for small sample data, in the feature space. (2) Similar to DELTA [13], we deploy regularizations to the feature extractor and the feature classifier for transfer learning to select the discriminative features from the outer layer outputs and transfer pre-trained knowledge from the large-scale source dataset. (3) In the training phase, we employ a feature generation method in a smallsized target dataset, which augments the selected discriminative features by adopting the attention mechanisms for each class label through supervised learning.
Our main contribution is a novel feature generation and augmentation using the proposed GAN structure to unravel the data imbalanced problem and improve neural network performance. In the proposed GAN structure, a feature extractor and a feature 2 classifier are included to train together with a feature generator and a feature discriminator. For the generation of meaningful features for classification of small-sized target data, transfer learning with regularization and class-wise attention is adopted. We evaluate the proposed method with various pairs of source and target dataset to show general applicability to classification in the class-imbalanced conditions. The remainder of this paper is organized as follows. In Section 2, related works are summarized. The framework and training procedure are presented in Section 3. In Section 4, the benchmark datasets for verification and the experimental results are detailed. Section 5 concludes the paper and provides scope for future research.

Related Work
GAN [4] is a generative model based on a min-max game theory scenario that pits two networks against one another. A generator network, G, competes against a discriminator network, D, that distinguishes between samples generated from G and samples from the training data. GAN can generate synthetic data close to the original; however, the training process of GAN has the instability of loss function convergence and the problem of mode collapse [16]. To avoid such problems, the Wasserstein GAN (WGAN) [17] and the WGAN with a gradient penalty (WGAN-GP) [18] using the Wasserstein-K distance as the loss function were proposed. Recently, GAN has been extended in several ways to control the generated properties [19,20,21] and its utilization, such as data augmentation [6,22,7].
Recently, GAN-based data augmentation methods were proposed to further improve the creation of augmented synthetic training data. Huang et al. [8] proposed stacked GAN (SGAN), which is trained to invert the hierarchical representations of a bottom-up discriminative network. Guo et al. [9] proposed a discriminative variational autoencoding adversarial network, which learns a latent two-component mixture distributor and alleviates the class imbalance for deep imbalanced learning. Cui et al. [23] proposed a class-balanced loss for long-tailed distributions. The class-balanced loss re-weights losses inversely with the effective number of samples per class. Although the generated synthetic data can balance the distribution between the classes, the classification performance is limited due to no guarantee enhancing the feature extraction ability of the classifier, when it is an independent network of the GAN. Even if GAN is trained with the classifier together like [24], GAN has the limitation to generate real-looking synthetic data on large-scale images. Not-qualified synthetic data is not able to improve the performance of the classification significantly.
Unlike GAN-based data augmentation at the input data level, Volpi et al. [15] proposed augmentation in the feature space. A feature extractor is trained with the source dataset under supervised learning, and then a feature generator for the unlabeled target dataset is trained in the convolutional GAN (CGAN) framework [19] against the feature extractor. Zhang et al. [14] developed a more general feature generation framework for imbalanced classification, inspired by the adversarial feature augmentation approach. These methods can generate domain-invariant features without considering the modality of the data. Consequently, they achieved better performance than unsupervised domain adaptation methods. However, the improvement of the classification under the classimbalanced condition is still not significant. The dimension of the features in the feature extractor is considerably lower than the input data. Such a low dimension may not be 3 sufficient to present the data distribution of a small amount of minority class data. The generated features should be domain-invariant and modality-free. Transfer learning is a machine learning method that transfers knowledge learned in a source task to a target task [25,26,27]. The weights of deep neural networks are first pre-trained on a large-scale dataset, which is called the source task, and then fine-tuned using the data from the target task with a small amount of data [26]. In the case of fine-tuning in convolutional neural networks, the weights in the first few convolution layers are fixed by pre-training. The last convolution layers are fine-tuned by the target task. During fine-tuning, the parameters of the target model can be driven far away from the pre-trained parameter values, leading to the incorporation of information relevant to the targeted problem and overfitting to the target task, so-called catastrophic forgetting [28]. This simple method cannot guarantee good performance because it may in many cases burden the network of the target task with irrelevant information. Yosinki et al. [29] quantified the transferability of features from each layer to learn transferable representations. For constraining catastrophic forgetting in inductive transfer learning, L 2 -norm regularization was proposed in [30]. The key concept of L 2 -SP is "starting point as reference" optimization, which tries to drive weight parameters to pre-trained values by regularizing the distance between the parameters of source and target tasks.
Instead of regularizing the weight parameters, deep learning transfer using a feature map with attention (DELTA), which is a regularized transfer learning framework, was proposed by Li et al. [13]. Inspired by knowledge distillation for model compression [31,32,33], DELTA employs the ideas of "inactivated channel re-usage" and feature map regularization with attention. These regularization approaches achieved significant improvement and alleviated the catastrophic forgetting problem by drawing weight parameters close to pre-trained values or aligning transferable channels in feature maps. It constrains the difference between the feature maps generated by the convolution layers of the source and target networks with attention. DELTA selects the discriminative features from the outer layer outputs using a supervised attention mechanism. To further improve the performance of the classification with a small training dataset or class-imbalanced dataset, a DFG method using attention maps in the feature space is proposed in this paper.

Discriminative Feature Generation
Our goal is to generate discriminative features to improve the performance of the neural network on a small number of training data and the class-imbalanced dataset.
Unlike CGAN and ACGAN [19,20], the structure of the feature classifier to the independent network is employed in the proposed GAN structure. The proposed GAN structure is composed of four independent networks: a feature generator, feature discriminator, feature extractor, and feature classifier. The structure of the proposed GAN model is shown in Figure 1. The feature extractor represents a first convolution layer or block in a neural network classifier, and the feature classifier represents the rest of the convolution layers and fully connected layers after the feature extractor in the neural network classifier. The feature extractor and classifier are regularized with transferred weights from the source model, similar to Li et al. [13]. We deploy the knowledge distillation technique to the training procedure of GAN and train the feature generator with supervised attention models to generate discriminative features. 4

Training GAN model with four independent networks
For the stability of the training procedure and the quality of the generated data, we deploy the objective formulation of WGAN-GP for the feature discriminator and the feature generator. The loss functions of the feature discriminator and the feature generator are denoted by L D and L G , respectively.
where z is the noise vector sampled from uniform distribution P z , and feature generator G generates synthetic feature G(z). λ is the penalty coefficient, and Px is the uniform sampling along straight lines between pairs of points from the real data distribution, P r , and the generated data distribution. θ D and θ G are the parameters of the feature discriminator and the feature generator, respectively. Feature discriminator D is trained to minimize L D to distinguish between real and generated features. Generator G is trained to minimize L G .
To generate the distribution of features similar to the real features from the feature extractor, a total loss function of the feature classifier contains a loss function of the 5 feature classifier on the features concatenating the real and generated features as well. The objective functions of the feature extractor and classifier are represented as follows: In (4), the loss functions of L r C , L g C , and L c C are defined as where L E and L C are the loss functions of the feature extractor and feature classifier, respectively. θ E and θ C are the parameters E and C, respectively. ⊕ denotes a concatenation operation. L c C is the loss function of the feature classifier on the feature maps concatenated with the real and generated features. α, β, and γ are hyperparameters that control the importance of the classification for the real and generated data. In L E and L r C , the regularization term, which characterizes the differences between the source and target network for transfer learning, can be added.
Feature generator G (to minimize (2) and (4c)) and feature classifier C (to minimize (4) and (4a)) are trained simultaneously, whereas feature discriminator D and feature extractor E are trained to minimize (1) and (3), respectively. The reason why feature classifier C is optimized not only with the real data is to prevent the classifier from overfitting to the real data because the performance of the classifier with the real data is not sufficient to improve the performance of the classification under the class-imbalanced condition.
In the training procedure, we define two training parameters. The generator learning parameter controls the ratio between L G and L g C while optimizing the generator parameter, θ G . The classifier learning parameter is used to control the balance of how much the classifier is trained from the feature extracted from the real data L r C or from the real features combined with the generated features, L C . The training details of the proposed method using (1), (2), (3), and (4) with these parameters are summarized in Algorithm 1.

Discriminative feature generation by using the supervised attention model
Algorithm 1 Training procedure of the proposed DFG method. We use the default values of λ = 10, n D = 5, n C1 = 2, n C2 = 10, andn W = 2000 Require: Batch size m, learning rate η, hyperparameter for weight sum ρ, hyperparameters α, β, and γ, and a threshold value for filter weight δ.

11:
Sample {x (i) } m i=1 ∼ P r a batch from the real data and labels {y (i) } m i=1 .

23:
for k = 1,..., n l number of class labels do 24: To generate discriminative features, we adopt supervised attention mechanisms for each class label. To obtain the weights for feature maps, we propose a supervised attention method adopted from [13] for the generator network. Whereas the supervised attention method in [13] is calculated by averaging the filter weight for each filter, we utilize the fact that the importance of each filter varies from class to class. We transform the filter weights for a single data into class-wise filter weights and deactivate channels with low filter weights.
In the supervised attention method from [13], the weights of the features are characterized by the performance loss when removing the convolutional filter for each feature from the feature extractor network. For a conv2d layer in the feature extractor and generator, the parameter form is a four-dimensional tensor with the shape of (c i+1 , c i , k h , k w ), where c i denotes the number of channels of the i-th layer and (k h , k w ) represents the size of the kernel. At the last convolution layer of the feature extractor, the dimension of the output feature maps is (c i+1 , h i+1 , w i+1 ), which is matched with the dimension of the output of the feature generator. We can measure significance through the performance reduction of the feature extractor and classifier when the filter in the last convolution layer is disabled in the feature extractor network. In other words, because it usually causes higher performance loss to remove a filter with a greater capacity for discrimination, we can improve the performance of the classification by generating feature maps focusing on channels with high filter weights. The filter weight is calculated through the feature extractor and classifier and is input into the feature generator. The feature generator generates weighted feature maps with deactivated feature maps using the filter weight. The deactivated feature maps are denoted in black in Figure 1.
The filter weight is expressed in (5), which is used to calculate the gap between the classification losses of the feature extractor and classifier on data (x, y) with and without the j-th filter in the last convolution layer of the feature extraction network.
where θ \j E denotes the modified parameter from θ E with all elements of the j-th filter set to zero.
Because each class usually has different importance of the filter channel, the calculated filter weights for a single data point can be transformed into class-wise filter weights.
where W j (E, c; θ E ) (W ∈ R nc×n f ) is a class-wise filter weight, c is a class label, n c is the number of images in the class, and n f is the number of filters. By multiplying the number of filters n f , we set the average value of the weights to 1. When transfer learning is used for the small training dataset and class imbalance problem, a filter weight on the source model can be obtained through feature extractor E S with parameter θ E S and classifier C S with parameter θ C S trained on the source dataset. Whereas the filter weight in [13] was calculated on the source model only, the proposed method updates the filter weight on the target model. In this case, the filter weight can be expressed in the form of 8 a weighted sum between the filter weight on the source dataset and the target dataset.
By applying the weight to the filters of the last layer in the feature generator, the feature generator can generate discriminative features by following the supervised attention models. The operation of the last layer in the original feature generator is expressed as follows.
Note that f l denotes the features of the l-th layer in the feature generator, norm() performs batch normalization, TransConv(f, θ) is a transpose convolution that takes f as the input with θ as a weight, and b l is a bias vector of the l-th layer. To generate discriminative features by using the class-wise weights of a filter, the operation can be rewritten as follows.
f l (z,ŷ) is a discriminative feature generated by the supervised attention model. To maximize the difference in the weights and inactivate unimportant filter channels, we set the weight values that are smaller than a threshold, δ, to zero and represent as W * j . δ is set to 0.95 of the average value of all filters for each class.
To evaluate the proposed method under the class-imbalanced condition, we adopt the step imbalance type described in [3]. The step imbalance assumes that the classes are divided into minority and majority groups. Then, classes in the same group have the same number of data points, making a step in the data distribution plot. In this evaluation, we set the number of majority classes as two and eight to evaluate under imbalanced conditions. All combinations of class imbalance cases were tested in the 9 evaluation. For small training datasets such as STL-10 and Caltech-256, we did not create an imbalanced condition. Under the class balanced condition, each experiment was repeated 10 times. SVHN: To test the transferability in the same domain, we set EMNIST as the source task and SVHN as the target task, where both are single-digit images. To reduce the image dimension gap, the SVHN images are converted to grayscale. For the imbalance 10:1 ratio set-up in training, the number of majority class images is 5000, and one of the minority class images is 500 in SVHN.
F-MNIST: To test the transferability in the different domains, we set EMNIST as the source task and F-MNIST as the target task which has images of fashion and clothing items. For imbalanced 40:1 ratio set-up in training, the number of majority class images is 6000 and that of minority class images is 150 in F-MNIST.
STL-10: STL-10 images are resized to 32 × 32 to match the image dimension of CIFAR-10. In this case, class imbalance is not applied, but a small training dataset where 500 images per class are used in STL-10.
CINIC-10: CINIC-10 dataset contains 9000 training images in each class and we sampled 900 images from each minor class for imbalance 10:1 ratio.
Caltech-256: Caltech 256 contains 257 object categories and 30607 images. In this study, we sampled 60 training samples for each category, referring to [30,13]. We resized the input images to 256 × 256, followed by data augmentation operations of random mirror and random crop to 224 × 224.
Food-101: Food-101 contains 101 food categories with 101000 images. A total of 750 training images and 250 test images were provided for each class. We sampled 250 training samples from minority classes for an imbalance of 5:1 ratio set-up during training.
All the experiments were implemented using Python scripts in the PyTorch framework and tested on a Linux system. Training procedures were performed on NVIDIA Tesla V100 GPUs. The training details can be found in Appendix A.

Results and comparisons
To establish a baseline for comparison, we denote "Original" as the performance of a classification model trained on the target dataset. "Original" is trained under the classimbalanced condition in the case of SVHN, F-MNIST, CINIC-10, and Food-101. In the case of STL-10 and Caltech-256, "Original" is trained under the class-balanced condition that contains a small number of train data. As mentioned, "Original" in Table 2 is trained without transfer learning and shows accuracies tested on the target test dataset achieved with classifiers trained on the target training dataset. For the class imbalance problem, in which classes are selected as majority classes, is an important factor in determining the performance of the neural network, and the deviation of result accuracies is large. Thus, accuracies under the class imbalanced condition show the mean and standard deviation.
The proposed method is compared with previous fine-tuning methods.
• Fine-tuning: Traditional transfer learning method by fixing weights before the last convolution block and fine-tuning the rest of them.
For a fair comparison with the other methods, we conducted experiments on the same neural network architecture for classification, such as LeNet5, VGGNet-16, and ResNet-50. Table 2 shows the results of the proposed DFG method compared to the finetuning, DELTA, and DIFA+cMWGAN methods. The results are the top-1 classification accuracies. The results in Table 2 show that our method based on DFG leads to accuracies higher than other methods for three different classifiers and six different benchmark datasets.

Qualitative results and ablation study
Although the proposed DFG method shows higher accuracy than the other methods, we would like to better understand the benefits brought by DFG and visualize how closely the proposed DFG generated the distribution of features resembled that of the original in two-dimensional spaces. We apply PCA [47] and t-SNE [48] analyses of real 11 and generated features. Figure 2 shows comparisons between the real and generated features. Different colors indicate different classes in the left part, and two different colors indicate real and generated features in the right part. From a qualitative point of view, real and generated features are distributed similarly, and the inter-class structure is preserved. Additionally, we analyzed the activation maps from the proposed DFG method using class activation maps [49]. From Figures 3 and 4, we observe that the regions of the target object have higher activation than that of the original classifier and fine-tuning. Compared with the activation maps from the DELTA, the proposed method improved the concentration and activation degrees.
We carried out an ablation study to better understand how the filter weight changes during the training procedure. We perform an experiment by changing the ρ value in (7). Figure 5 presents the distribution of the weights of two representative filters in the case of the SVHN dataset for two different hyperparameters: ρ = 0 and 0.75. The first row in Figure 5 represents the distribution of the weights of two representative filters with ρ = 0.75, and the second row represents the distribution of the weights of the same filters with ρ = 0. The red straight line is the line of equality for all filters. For instance, (a) if a weight value is placed above the line of equality, it means that the importance of the filter for the class is higher than the average of all filters. The case with ρ = 0 involves the training of the feature generator with the weights of the filters by updating from the feature extractor trained on the target dataset, and the case with ρ = 0.75 involves the training of the feature generator with the weights of the filters by two feature extractors: one is trained on the source dataset and the other is trained on the target dataset in the training procedure. We observed a significant difference between the two distributions of the weights. The distribution of the weights of the filters with ρ = 0 converges to the equality baseline, whereas the distribution of the weights of the filters with ρ = 0.75 maintains the characteristics of the weight for each class label. Our hypothesis of a possible cause of the difference is that the feature extractor trained on the target dataset tends to be trained to activate all filters and be overfitted to the target dataset. The feature extractor trained on the source dataset inactivates several convolution filters, and a small number of filters are activated. The average accuracy with ρ = 0.75 is 80.81% and that with ρ = 0 is 79.98%. The weights of the filters with ρ = 0.75, using the weight values from the feature extractor trained on the source dataset, prevent overfitting and help to generate discriminative features for each class label. Figure 6 shows the effect of hyperparameter ρ on the classification accuracy on the SVHN dataset using DFG. In this case, DFG with hyperparameter ρ = 0.75 outperformed that with other hyperparameters including the cases that the weights of the filters by the feature extractor only trained on the source or the target dataset. The test accuracy varies smoothly according to hyperparameter ρ, and combination of the two different weights of a filter yields a better accuracy. 15  Figure 6: Classification accuracy on SVHN using DFG by changing hyperparameter ρ.

Conclusions
In this paper, we proposed a novel DFG method using attention maps in the feature space. The structure of the feature classifier network of the independent network was employed in the proposed GAN structure. To generate discriminative features, we adopted supervised attention mechanisms for each class label and applied the mechanisms to the feature generator network for the filter weights and feature activation. The experimental results showed that the proposed DFG method improved the accuracy of the classification models of LeNet-5, VGG-16, and ResNet-50 by 4.24%, 5.52%, 14.73%, 13.95%, 42.99%, and 45.83% on six different benchmark datasets, respectively. Comparing to the transfer learning method, DELTA, and the other adversarial feature augmentation methods, the proposed DFG method provided the largest improvement in classification performance. Furthermore, we visualized the distribution of generated features by applying PCA and t-SNE and analyzed the activation maps from the proposed method using class activation maps. The qualitative results showed that the features generated by the proposed method are distributed similarly to the real features and the proposed method improved the concentration and activation degrees in the class activation maps.
In the future, we will extend the proposed method to real-world problems and contemplate further accuracy improvement in the feature generation method. Although the proposed method improved the performance of the classification under the classimbalanced conditions, we believe there is room for improvement as a gap still exists between our method and the balanced condition.