{'fraction_non_alphanumeric': 0.03235676333667813, 'fraction_numerical': 0.015331216791959337, 'mean_word_length': 4.410947930574099, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Attention-based Convolutional Autoencoders for 3D-Variational Data Assimilation

We propose a new 'Bi-Reduced Space' approach to solving 3D Variational Data Assimilation using Convolutional Autoencoders. We prove that our approach has the same solution as previous methods but has significantly lower computational complexity; in other words, we reduce the computational cost without affecting the data assimilation accuracy. We tested the new method with data from a real-world application: a pollution model of a site in Elephant and Castle, London and found that we could reduce the size of the background covariance matrix representation by O(10^3) and, at the same time, increase our data assimilation accuracy with respect to existing reduced space methods.


Introduction
Data Assimilation (DA) is an uncertainty quantification technique in which observation data and a forecasting model are used in tandem to generate predictions that are more accurate than those that would be produced using either component independently. DA is computationally costly for large systems [1] and under operational constraints, it is often necessary to solve the problem in a reduced space in order to achieve real-time assimilation. In most relevant DA operational software, a variable transformation is performed on the variational functional to reduce the computational cost needed for computing the covariance matrix explicitly; to reduce the space, only Empirical Orthogonal Functions (EOFs) of the first largest eigenvalues of the error covariance matrix are considered. Since its introduction to meteorology by Edward Lorenz, EOFs analysis, which is essentially based on a TSVD, has become a fundamental tool in computational fluid dynamic modelling for data diagnostics and dynamical model reduction. Real world applications of TSVD (EOFs) basically exploit the fact that these methods allow a decomposition of a data function into a set of orthogonal functions, which are designed so that only a few of these functions are needed in lower-dimensional approximations. Nevertheless, the accuracy of the solution obtained by truncating, exhibits a high sensitivity to the variation of the value of the truncation parameter [2,3], so that a suitable truncation parameter is needed. This is a severe drawback of truncation-based methods and limits the utility of operational software based on these methods. Convolutional Autoencoders (CAEs) have had huge successes in computer vision [4,5] and particularly in image compression [6,7]. In this work, we use CAEs to produce a reduced space in which DA can be performed efficiently.
The structure of this paper is as follows: in Section 2 we cover related work and we present the contribution of the present work. Section 3 provides preliminary concepts and definitions, and Section 4 introduces our theoretical contribution. As the success of our approach is heavily conditioned on the choice of a CAE architecture, in Section 5

Preliminary Definitions
In this section we define the key quantities for DA. In most cases, we follow the notation in Banister's review paper [24].
• Let x t represent the state of the model at time t such that: x t ∈ R n (1) where n is the number of elements in the model state vector. The state for T time-steps can be given in a single matrix: In most practical problems, n is large and of order ≥ O(10 6 ). • Let y t represent the observation space of the system where: where typically M << n. The Met Office uses M = 0.01n [1]. • Let H t be an observation operator such that: where t ∼ N (0, R t ) is the observation error. Often, observations are assumed to be uncorrelated meaning that R t is diagonal. When all observations are of the same type and made with the same device we have: • Let M t−1,t be the forecast model that propagates the system forward from time-step t − 1 to t such that: where η t ∼ N (0, Q t ) is the model error introduced over this interval.
• Let x b t ∈ R n be the background state at time-step t. All a priori information about the system is introduced through the first background state x b 0 and the model M t−1,t . Future background estimates are then defined according to the free-running model in which η t is assumed to be zero and therefore: • Let B t represent the background state x b t covariance matrix. In theory, it is found by evaluating: where x * t is the true state of the atmosphere. In practice, even if x * t were known, this matrix is too large to fit in memory as it has O(n 2 ) parameters which is ≥ O(10 12 ) for most practical problems.

Variational DA, VarDA
VarDA involves minimising a cost function in order to find the most likely state values X DA given the observations y t , the model predictions and their uncertainties. The problem is to find the initial state x DA 0 that satisfies: J(x 0 ) = 1 2 The first term in cost function J(x 0 ) measures the difference between the initial model state and our a priori expectation of this state. The second term encodes the difference between the observations and the model forecasts. The cost function is explicitly differentiated and then approximately solved by first-order minimisation routines. Note that we are assuming all errors are Gaussian by using this least-squares formulation 4 . This formulation is known as the 'strong constraint 4D-Var' where: • '4D' refers to the fact that we are considering three spatial dimensions as well as one temporal dimension. It is contrasted with 3D-Var in which a single time-step is assimilated. • 'Strong' refers to the fact that model errors are assumed to be zero [?] 5 .
In the current work we are considering the 3D case where the cost function is: Note that we have dropped the t subscripts as we are only assimilating a single time-step.

Incremental VarDA
If H t and M t−1,t are linear (which they are not in general), the cost functions (10) and (11) are convex. We can approximately linearize these operators about the background state x b by formulating the problem in terms of perturbations to this state in a method known as the incremental formulation [11]: The problem statement then becomes: where H is the observation operator linearized about the background state and the 'misfit' between observation and expected observation is:

Control Variable Transform
As B is in O(n 2 ) it must be represented implicitly. A common way of doing this is by using the formulation proposed in [19] which states that: where V is the Cholesky factorisation of B. In this case the problem can be written as: Following (8), we can see that V is theoretically found by stacking a series of background states X b and subtracting the true state, x * : In reality, we do not know x * but we can estimate V with a sample of S model state forecasts X b that we set aside as 'background' such that: where x b :=x b is the mean of the sample of background states. In the incremental 3D-VarDA formulation this gives: where Note that w ∈ R S and since S << n in all practical cases, w is referred to as the 'reduced space'. V is an affine transform from the reduced space to the full space and is used to obtain the assimilated state after minimising J(w): The problem in (17) is poorly conditioned because, in most practical contexts, the matrix V is poorly conditioned 6 .

Truncated SVD
One way to precondition V and increase the speed of convergence is to use an eigenanalysis technique such as SVD [26,27] or PCA to generate Empirical Orthogonal Functions (EOFs) [28] and remove low-variance modes from V . We have implemented 3D-VarDA with TSVD as described in [17] and we evaluate the success of our proposed approach against this scheme in Section 6. In order to draw out the theoretical differences between the methods we briefly 6 In other words, the ratio of V 's largest to smallest eigenvalue is large.
summarise TSVD here. The singular value decomposition of matrix V is as follows: where U ∈ R n×S , Σ ∈ R S×S are both orthogonal and W ∈ R n×S is diagonal and contains V 's eigenvalues σ i such that: where the eigenvalues appear in decreasing order: To perform TSVD, τ of the modes are retained where 0 < τ < S and the matrix V τ is reconstructed such that: This has generalised inverse: where Σ + τ is the generalised inverse of Σ τ such that:

Autoencoders
AEs are a self-supervised machine learning method first proposed in [29]. They consist of two components: an encoder f (x) which compresses the input x to a small representation z (the AE's reduced space is referred to as the 'latent space' in the machine learning literature), and a decoder g(z) which reconstructs the input: where the input data x and reconstructionx are both ∈ R n and the latent representation z ∈ R m . As m < n, the 'information bottleneck' forces the AE to find and exploit redundancies in the training data in order to implicitly model the data distribution 7 . PCA is a special case of an AE with linear activations [31].
Typically AEs are trained with the L1 or L2 reconstruction error. In this work we found that networks trained with the L2 loss I 2 (x,x) performed consistently better for DA than those trained with I 1 (x,x). We experimented with fine-tuning our models with the L1 loss as suggested in [32] but this did not provide a consistent improvement.
AEs are useful in any problems in which a latent data representation is required for a downstream task but they have had success as a standalone solution in anomaly detection [33,34] and image denoising [4,5]. There are many variants of Autoencoder but we have focused on the image compression CAE literature to guide our AE design. Variational AutoEncoders [35] have achieved recent success in a range of generative modelling tasks, particularly NLP [36,37], and enforce the latent orthogonality condition that we will discuss in Section 4 but they typically produce poorer reconstructions than vanilla AEs (see chapter 20 of [38]) so we do not use them in this work. Similarly, Huang et al. proposed a GAN-based AEs to generate visually plausible reconstructions [39] but there is a lot of work to be done in proving the 'correctness' of GAN-generated samples and we were not confident that 'visually plausible' in the image domain would translate to scientifically correct in the 3D field domain. Nevertheless, this is certainly another route for future work.
The majority of CAE architectures in the DA literature are in the field of Reduced Order Modelling (ROM) [40,41,42,43]. ROMs are only applicable in 4D-VarDA when an online method of generating model forecasts is required.
Their architecture reflects this: as this is a sequence-to-sequence problem [44] most modern ROMs use LSTMs or other RNN variants to make their predictions. The result is that these architectures are not applicable in this work. Instead, we have reviewed the image compression literature and implemented a range of state of the art networks from this field. In the remainder of this section we describe the image compression problem and the SOTA networks for this task. The details and results of our architecture search are deferred to Section 5.   and while success on one metric generally implies success on another, this is not always the case. When there is ambiguity, we have given the highest positioning a system achieved. In the final column we have provided the name of the system according to the naming conventions we define in Section 5.1 -we implemented and evaluated all of the named variants. In this column a '∼' symbol implies our implemented system is similar, but not equivalent, to the model in question. * To our knowledge, these teams did not produce a publication detailing their approach. † This system was also the fastest in 2019. § The original model did not have skip-connections but it otherwise identical to our 'ResNeXt3-3-N-vanilla'. ‡ The authors use 'wide-activated residual blocks' in which the channel size is increased by a factor of 4 internally within the block. The ResNeXt authors [23] have shown that this is equivalent to their system with a cardinality of 4. ∇ This system was not an end-to-end AE but used the VVC compression standard [60].

AEs for Image Compression
In lossy Image Compression (IC), an image is contracted to a more space-efficient representation with a small loss of information. Since 2017, CAE-based IC systems have started to outperform traditional IC algorithms [6] such as JPEG and JPEG 2000 8 . There are a few key differences between the IC problem and ours 9 but the similarities mean that the IC-CAE literature provides a useful starting point. Crucially, unlike physical-field compression, the IC problem has received a great deal of attention from the machine learning community. For example, CLIC or 'Challenges on Learned Image Compression' runs annually during CVPR to find the state of the art in lossy IC. We used the papers of the CLIC winners and runners up for the years 2018 [54,55,56] and 2019 [20,56,57,58,32,59] and followed their citations to give approximately 40 relevant machine learning papers. The successful entrants used a variety of architectural components including attention-based models [48], GDNs (Generalised Divisive Normalisation transformations) [45], multi-scale learning [47] and complex residual blocks (RBs) [?, 23, ?, 50, 51] in order to improve their compression quality. A summary of these CAE architectures and an overview of their constituent elements is given in Table 1. We have implemented a variant on all of these systems according to a framework described in Section 5.1.
Of particular note, are the Tucodec team who came first in CLIC-2018 and CLIC-2019. In 2019, two of the top-five finishing teams (Joint [58] and Vimicro [62]) used a network that was virtually identical to Tucodec's entry from the previous year. However, the Tucodec team improved on their previous design with the addition of Residual Attention Blocks (RABs) 10 and won the competition for a second year running. In our experiments in Section 5.3 we found that our 3D extension of their 2019 system (see Figure 1) performed consistently better than the other architectures that we considered. 7 Although they are rarely thought of in this way, AEs are theoretically equivalent to a clustering algorithm [30] in the sense that the encoder learns to map commonly co-occurring inputs to a single internal representation. 8 In fact, a new lossless CAE-based IC format L3C was proposed in May 2019 that also outperforms traditional lossless methods [61]. 9 Specifically, that the encoder output in an IC system is a bitstream while in our case it is simply a vector of floats. This adds complexity in comparison with our system in a number of ways: firstly, there is a tradeoff in IC between the compressed size and quality so most systems are trained with a multi-task loss of bitstream entropy and reconstruction error; secondly, the quantization operation is non-differentiable and therefore the systems cannot be trained by backpropagation directly; and finally these systems require 'importance maps' [62], a form of attention to determine how many bits should be allocated to each region of the image. 10 We note that the Tucodec authors assert incorrectly throughout their paper that they are using non-local attention blocks or 'RN-ABs' instead of the simpler local-attention RABs (proposed in the same paper by [51]). We have used RABs in our implementation of their system. Figure 1: The Tucodec encoder [20]. The decoder is symmetric except that it does not have the multi-scale path. All parameters are specific to our 3D implementation.

'Bi-reduced space' formulation of DA
In this section we describe our proposed DA formulation with AEs and discus its computational cost. The new approach involves non-trivial changes to the incremental CVT formulation in (17). A subtlety to note is that there are two 'reduced' spaces in this case: the reduced space of size S introduced by the CVT and the reduced space of size m introduced by the encoder-decoder framework. As our method utilises both of these spaces we refer to it as the 'bi-reduced space' formulation. To avoid confusion, we describe the AE space as the 'latent space' (denoted with z) and the CVT space as the 'reduced space' (denoted with w).
In Section 4.1, we define a number of quantities required for our formulation in 4.2 and introduce a series of constraints and assumptions under-which we can show (in Section 4.3) that the proposed approach is equivalent to the Parish et al. CVT formulation [19].

Constraint: mean-centred data
All data is mean-centred with respect to the historical meanx b . As this sets x b = 0 for all equations in Section 2, (29) has multiple equivalent expressions: This constraint is necessary for Section 4.3 but has the added benefit of ensuring that encoder inputs are symmetrically distributed about 0. The absolute state can be reconstructed by addingx b to the mean-centred state. Whenever the DA performance is evaluated in this paper the absolute state is used.

Definitions
We define a series of quantities below where a subscript l implies the matrix or vector has been replaced by its latentspace equivalent: • Let f o (·) be the 'observation encoder' which maps from the full observation space of size M to the reduced observation space of size M l such that: where H + ∈ R n×M is the generalised inverse of H which maps from the observation space to the full space and H l f H + (·) is the compound operator where the three transforms are applied in the order H + , f, H l to an observation ∈ R M . This is an under-determined problem so there will be many equivalent H + operators. H l ∈ R M l ×m is the latent space observation operator which maps from the latent space of size m to the latent observation space of size M l . We have implicitly defined H l := I which implies M l = m. Another way of thinking of this is that the entire latent space is observable to us.
• Let V l be the latent space equivalent of V such that: Note that we are representing the information in matrix V of size R n×S in a matrix of size R m×S . In our implementation m ∼ 0.001n so this is an O(10 3 ) reduction in the size of our background covariance representation. • Let d l be the latent misfit such that: • Let R l be the latent observation error covariance. Recall that the full-space covariance R is computed over the observation error = Hx b − y and as such, for observations y, can be calculated as R = E[ T ]. We define an equivalent observation covariance matrix R such that:

Assumption: Orthonormal latent features
If the full-space observation errors are uncorrelated (and R = σ 2 0 I), then the latent observation errors l will also be uncorrelated as they are derived directly from these full space errors giving: where σ l is the latent observation error standard deviation. However, we note that the uncorrelated observation assumption breaks down as the observation locations move closer to one another (i.e. as M increases). In this work we consider the case in which M = n, pushing this traditional DA assumption to breaking point. Nevertheless, (37) will hold in all scenarios if we impose the additional constraint that all latent features are orthonormal to one another. It is worth making a brief comment on the conditions under which this will be true: 1. Latent dimensions are not orthogonal in the general case but they will be 'approximately' perpendicular. To see why intuitively: if an AE is producing a good reconstruction, the latent hyperplane must span the majority of the data distribution manifold. If it is able to do this with a very small number of latent features then it must be rare for latent features to double-up and span identical areas of the data distribution. 2. It is possible to enforce latent orthogonal features with the use of a Variational Autoencoder [63]. This approach is not explored in this paper for reasons discussed in Section 5 but this is a clear avenue for future work. 3. To obtain unit length features, batch normalisation could be used in the encoder. However, like Chen et al. [55], we found that batch normalisation greatly hampered the AEs ability to produce good reconstructions and hence we did not use it in the backbone network of our models. For this reason, we do not make statements on the relation between the magnitude of σ 0 and σ l (although in practice we treat them as equal).

Observation encoder in practice
In the general case, calculation of the latent misfit d l in (34) requires evaluation of the observation encoder operator f o on the full-observation space: where as a result of the mean-centering constraint we have: We hypothesise that f o (y) could be modelled with a CAE framework that adds an implicit interpolation network H + to the trained encoder network f (·) as shown in Figure 2. This system could be trained with the reconstruction error over y andŷ = Hf o (y).
However, in this work we are using synthetic data meaning that the full state is available as observation. As such, we sidestep this complexity and use: which implies: In other words, we use the encoder to obtain our latent misfit. Verifying the feasibility of the scheme in Figure 2 is a necessary criterion for using the proposed approach operationally.

Proposed 3D-VarDA formulation
Our proposed bi-reduced space formulation is: Once this has been minimised in the reduced space the result can be restored to the full space in a two-stage transformation: 1. Multiplication by V l to move from the reduced space representation ∈ R S to the latent space ∈ R m . 2. Applying the decoder g(·) to move from the latent space to the full space ∈ R n .
Overall this gives: Note that, allowing for the different definitions of the latent variables V l , d l and R l , the formulation in (43) is identical to the mono-reduced space formulation in (17). As such, we were able to use the same cost-function, gradient and minimisation implementations when comparing the two approaches 11 .

Proof of equivalence
Our proposed formulation is equivalent to the mono-space formulation in (17) in that: This is true under three assumptions: 1. For high-performing autoencoder, we assume that the AE compression is lossless meaning g(f (x)) = x.
2. All features in the latent representation z are orthonormal as discussed above.
3. The observation space contains sufficient information to construct a good approximation of the full space x. This is a necessary condition in the creation of the f o operator in Figure 2 and is more likely to hold when M is large. This assumption is discussed in more detail in the proof of Lemma 5 below.
We will state and prove a series of lemmas to produce the result in (45).

Lemma 1
Let w DA and w DA l denote the solutions of (17) and (43) respectively, we have that Proof: The gradient of (17) is: Setting this to zero and solving for w will give the optimal value w DA as required to complete the proof: As the cost functions in (17) and (43) are mathematically equivalent where each operator is replaced by its latent equivalent in the bi-reduced space formulation and H l = I, we can write the exact solutions of w DA and w DA l in the following form: where, the matrices A and A l are such that:

Lemma 2
Let A and A l be the matrices defined in (48) and in (49) respectively. The following result held:

Lemma 3
Let b and b l be the matrices defined in (50) and in (51) respectively. The following result held: The overall result in equation (45) follows directly from Lemmas 2 and 3 but to prove these we need two further results in Lemmas 4 and 5.

Lemma 4
Let V l be the reduced matrix as defined in (33), f o defined in (32), H as in (14) and V as in (20). The followind result held: Proof of Lemma 4: by definition, we have then, from (55) we have V l ≈ f H + HV (56) Then the thesis follows from the definition of f o in (32): We note that H + H acts is an information bottleneck operator in which only information contained in the observation locations is propagated from V . Here we can assume that f V ≈ f H + HV due to some weak assumptions: • The observation space of size M contains sufficient information to construct a good approximation of the latent representation. This is equivalent to assumption iii) above because, if the observation space contains all information in the full space, by assumption i) it should also contain all information in the latent space. To show why this might be true consider that there must be large redundancies in the full space in order for the CAE framework to have any success. We argue in the following section that, in all practical scenarios, m, S < M . If a state of size m can contain most of the information of a state of size n it is not unlikely that a state of size M might contain the same information. More concretely, in Section 6 we demonstrate that the Arcucci et al. CVT with TSVD DA method [17] suffers no degradation in accuracy when just 10% of the total state space is used as observations (M = 0.1n) and there is only a 5% degradation when M = 0.01n. • The reduced space of size S also contains sufficient information to construct the latent representation of size m. This condition is implied by the lossless compression assumption i) as the reconstruction of the full state passes through the reduced space and the latent space. Lemma 5 Let R l be the reduced matrix as defined in (37), f o defined in (32), H as in (14) and R as in (5). The following result held: Proof of Lemma 5: because the observations are uncorrelated, we have: Proof of Lemma 2: we use Lemmas 5 and 6 to give: from (57) and (54), we have: which gives the (52).
This completes the proof that w DA = w DA l .

Advantages over TSVD: Theory
Eigenanalysis techniques such as PCA and TSVD are alternative methods of producing reduced space representations of data and, as discussed, have been used canonically in preconditioned 3D-VarDA [17]. Having summarised the key components of our proposed system, it is now possible to discus the theoretical reasons why our method produces a) higher quality compression and b) is faster than the traditional methods. We verify these advantages experimentally in Section 6.

Compression quality
A well-trained CAE will produce higher quality reconstructions than those using TSVD for a number of reasons: 1. The mean of the training data distribution can be stored 'for free' in the decoder leaving space in the latent representation to encode sample variations.
2. CAEs explicitly use location data and can therefore utilise properties like local smoothness in order to compress the input more efficiently 12 .
3. The latent features are created from non-linear combinations of the inputs meaning they are likely to be of greater expressive quality. This eigenanalysis approach is only optimal when the data is drawn from a Gaussian distribution 13 .
4. By design, in truncated SVD, some of the information is intentionally discarded. This is not the case in the CAE framework.

Computational Complexity
The proposed method also has lower computational complexity in an online setting. Here 'online costs' refers to any calculation that must take place when a new set of observations are made. 'Offline costs' are everything else and includes the TSVD computation and the CAE training. The online complexities of the Parish et al. reduced space approach R on [19] and our bi-reduced space approach B on are: where I 1 is the number of iterations in the reduced space VarDA minimisation routine, M is the number of observations, S is the reduced space size (which is equal to the size of the historical data sample) and m is the latent dimension size in our proposed method. We note that to achieve comparable accuracy with the two methods we will typically have S > m, (and M > S) so R on > B on . We derive these results in the following sections.
We note that our proposed method's online complexity B on is independent of the number of observations M meaning it is never necessary to arbitrarily reduce the number of assimilated observations in order to meet the latency requirements of the system. We also note that the cost of training a CAE is considerably larger than of performing TSVD but as these operations occur offline they are not of primary importance in the creation of an operational system. We give the derivation of the online and offline complexities in the following sections but first discuss the encoder and decoder inference complexities: f (x) and g(z) complexity With input of size n and output of size m, the encoder and decoder inference complexities are of order O(nm).
Justification: in the simplest possible encoder consisting of a single fully connected layer, the complexity of mapping from the full to the latent space would be O(nm) exactly. The convolutional case is more complex and will be given by O(nK) for a architecture-specific constant K but we think that logical CAE design choices give m ≈ K since: 1. K will not be << m as this would mean the encoder was introducing the information bottleneck at a location other than the latent space. 2. Similarly, K should not be >> m as this would negate the computational benefit of using convolutions over a linear network.
A symmetric argument gives the decoder complexity as O(nm).

Online Computational Cost Derivation
There are two steps that contribute to the online-cost: 1. The evaluation of the cost function and its gradient during the minimisation.
2. Restoring the calculated w DA to the full space.
In the following, we use the symbol · to indicate the operation under consideration. 12 Note that DA localisation approaches do utilise location information but not in a way that increases compression quality [64]. 13 More formally, PCA truncated at mode τ gives optimal reconstruction for all linear models of rank τ .

Lemma 6
The online complexity of step i) in the Reduced space method is: where I 1 is the number of iterations in the minimisation routine.
Proof of Lemma 7: we repeat the cost function (17) here for convenience: A naive implementation of 64 (and its derivative) would be dominated by the matrix multiplication H ·V but, as H and V always appear together, this quantity can be pre-computed and the minimisation complexity is independent of n. A single iteration of the VarDA minimisation has complexity O(M S + M 2 ) where the first term originates from HV · w while the second is from where we are assuming R is diagonal. With I 1 iterations this gives Lemma 7.

Lemma 7
The online complexity of step i) in our bi-reduced space method is: where I 2 is the number of iterations in the minimisation routine.
Proof of Lemma 8: this argument is almost identical to Lemma 7 except that we replace HV ∈ R M ×S with V l ∈ R m×S in the cost function. Altering these dimensions gives the required result.

Lemma 8
Restoring the calculated w DA to the full space in the mono-reduced space formulation is in O(nS).
Proof of Lemma 9: the product V · w DA is in O(nS).

Lemma 9
Restoring w DA l to the full space in the bi-reduced space formulation is in O(nS).
Proof of Lemma 10: this requires computing V l · w DA l followed by g · V l w DA l which has complexity O(mS + nm) = O(nm).
This gives an an overall reduced space online complexity of: and a bi-reduced space online complexity of: Hence, the key comparison between the online complexity of the two methods is down to the relative sizes of variables I 1 , I 2 , M, n, m and S. We assert that in most practical cases: As many of these are user-chosen parameters we cannot prove that (67) always holds but we can make concrete arguments about their ranges in practical settings. Argument 1: I 1 > I 2 in the majority of cases.
Justification: We expect the problem to be better conditioned in the bi-reduced space than in the reduced space for the same reasons that it is better conditioned in the reduced space in comparison with the full space [65]. . They employ a combination of KFs and VarDA approaches [66] but in the VarDA scheme, it is implausible that they would use anything close to S = 10 7 as V would be a matrix of size 10 7 · 10 9 = 10 16 . We believe the same constraints will hold in all practical scenarios even if the number of observations is relatively small.

Argument 3: S > m for useful systems.
Justification: We found that a value of m that was a factor of x2.5 smaller than S gave superior DA accuracy compared with the traditional method.

Argument 4: B on = O(nS) in the vast majority of cases.
Justification: According to Argument 3 we have: We found I 2 = O(10) meaning I 2 S << n and hence B on = O(nm). Concretely, with our values of m = 288 and S = 791 we found that upwards of 97% of the execution time of the bi-reduced DA procedure was restoring w DA l to the full space.
Combining Arguments 2-4, gives the stated results in the previous section: and since we have argued that n < S, this implies R on is strictly greater than L on . When M is large (M ≥ O(0.05n)) we can go further than this: Justification: Clearly the exact values here will vary from one problem to another but we found that with M = (0.01n), steps i) and ii) had approximately equal execution time but when M rose much above this, the minimisation term dominated as a result of the quadratic complexity.
We investigate the negative effect on DA accuracy of using fewer observations in Section 6.

Computational Complexity: Offline
In comparison, the offline costs are much larger for the proposed method in comparison with the formulation with TSVD: Reduced space offline cost R off = O(nS(M + S)) (68) Bi-reduced space offline cost B off = O(nSEm) (69) for E epochs of training. We prove these results below.
Typically, Em > M, S so R off < B off . In practice we found that R off << B off : training a model took something on the order of 10-20 hours on a GPU whilst TSVD required approximately 10 minutes on the CPU.
Derivation of (68): calculating SVD for V ∈ R n×S where S < n has complexity of O(S 2 n) [67]. It is also necessary to precompute H · V τ which is in O(M nS) = giving the overall complexity in (68).
Derivation of (69): training for E epochs, with S historical samples, requires ES evaluations of the encoder/decoder giving a complexity of O(SEmn). We must also precompute V l = f · V which is in O(nmS) but this is comparatively small so the offline cost is as in equation (69).

Architecture Search
In this section we describe our architecture search framework (Section 5.1) and training configuration (Section 5.2) before detailing the results of this search in Section 5.3. We found that the proposed approach was only successful relative to the Parish et al. approach when state of the art CAE architectures were used. This point is worth highlighting as many Data Assimilation practitioners use very simple neural networks in their research [68,69].

Search Framework
In our process of finding a good CAE architecture, we were concerned that the successes of 2D image compression architectures might not be transferable to 3D spatial inputs. We attempted to minimise this risk by designing a framework within which it was feasible to systematically search an architecture space that approximately encompassed the design of every 2018/19 top-5 CLIC finisher in Table 1. In order to achieve this, it was necessary to make small alterations to some of the original systems. As such, although we found that the Tucodec model was vastly superior to our implementations of the other systems, it is possible that our small design variations mean the quoted performances are not representative of the original designs. We were not unduly worried by this possibility as our aim in this process was to find systems that performed well in our domain rather than make exact comparisons between image compression networks.
Following best-practices from the literature and the results from our early-stage experiments, all CAE designs that we searched had the following features: 1. Encoder down-sampling occurs via strided-convolution [70] rather than max pooling [71] so that the network can learn its own sub-sampling routine.
2. Encoder and decoder layers are mirrored to reduce the design space size 14 . Strided convolutions in f (x) are replaced with transposed convolutions in g(z).
3. Convolutions have kernel size k = 3 or less to reduce the number of parameters and computational cost (as these scale with k 3 for a 3D feature-map [46,72]).
4. Convolutional down-sampling parameters are generated by our ConvScheduler class that has priorities of: (a) Avoiding addition of padding in later encoder layers as these can introduce artefacts in the reconstructed state. (b) Avoiding creation of feature maps that are not centred on the input as these are much harder to reconstruct in the decoder. This is achieved by refusing stride, s, padding, p and kernel size, k combinations that, when acting on input of width W , result in rounding in the floor operation when calculating the output size = W −k+2p s + 1.

Batch
Normalisation is not used in the down-sampling backbone of the CAE architectures because our preliminary investigation showed that it resulted in reconstructions of considerably poorer quality 15 . This finding was also made by Chen et al. [55].
6. Batch Normalisation is used in the residual blocks of the network. This was necessary to prevent activation and gradient overflow in deep networks.
7. Latent size is fixed at m = 288 to enable like-with-like comparisons between CAEs. This value was a reduction in state size by approximately three orders of magnitude.
For all models we investigated: 1. Three activation functions: ReLU [73], GDN [45] and PReLU activations [74]. Figure 3) each with and without the lightweight CBAM module [50] after the original RB. Figure 3: The two basic RBs we evaluated. We also investigated the effect of placing CBAM blocks after each of these (see A for a diagram of the CBAM RB).

Four RBs: vanilla [?] and NeXt RBs [23] (see
Note that vanilla RBs have considerably more parameters than NeXt RBs: for the 32 channel input versions shown in Figure 3, NeXt blocks have 2k parameters while vanilla blocks have almost 60k parameters. We found that we were able to reduce the Tucodec decoder inference latency by x2.5 by replacing vanilla RBs with NeXt RBs.
In designing a search strategy we noted that the models in Table 1 fit into one of two categories: 1. They are closely based on the Tucodec 2018 entry.
2. The encoder alternates between residual feature extraction and fully convolutional down-sampling operations.
We found that we could capture most of the variation in the second category with the 'backbone' encoder architecture shown in Figure 4. This design, without any 'optional blocks' is a fully convolutional network with seven layers, the latter five of which down-sample the feature map. It is very similar to the encoder of Theis et al. [6] that was used in the first system which outperformed JPEG compression. The backbone is responsible for down-sampling while the optional blocks can introduce innovative feature extraction mechanisms. We require that any added blocks do not change the feature map size and are bypassed with a skip connection so that, at least in theory, they cannot hinder the down-sampling process. In Section 5.3 we evaluate four variants on this backbone that are summarised in Table 2. For more information on these systems, including a discussion of our naming conventions, see A.   Table 2: Variants on the backbone encoder in Figure 4 that are evaluated in Section 5.3.

Data
We used simulated data from a single run (988 time-steps) of the open-source, finite-element, fluid dynamic software Fluidity on a small domain in Elephant and Castle in South London. This system had 100,040 states spread over a region of size (x, y, z) = (700m, 650m, 250m). Fluidity uses an adaptive unstructured mesh in order to provide high resolution in regions of interest without requiring this same resolution in locations where there is little variation. As convolutional kernels work on the assumption that adjacent states are equally spaced, it was necessary to interpolate between the points in the unstructured mesh to create a regular 3D grid with 247,520 evenly spaced points in the shape (91, 85,32) 16 . When evaluating the reduced space approach, this interpolated input was flattened before use. The data was placed in time-step order and the first 80% was used for the training set. All of the data was normalised using the training data statistics. In order to ensure a fair comparison with traditional TSVD methods, the reduced space V τ was calculated using the training set only.

Augmentation
With a view to increasing model generalisation, we investigated a number of regularisation techniques. We found that dropout harmed performance, even when it was applied channel-wise, and only to latter layers as recommended in [75]. Similarly, preliminary experiments showed that weight decay resulted in a small degradation in performance. As such, data augmentation was our only method of regularisation. We did not find any augmentation strategies for physical fields in the literature and decided that the only appropriate augmentation strategy was one borrowed from imaging: 3D 'field-jitter' (the 3D mono-channel equivalent of colour-jitter). This involves injection of normal noise of amplitude pσ at r of the state locations where p, r ∈ [0, 1] and σ is the state standard deviation. With regards to other image augmentation strategies, we did not judge it appropriate to crop the inputs because the state space is of fixed size and hence there is no benefit in the CAE learning to compress inputs at variable sizes. Similarly, we did not flip the inputs horizontally as there are buildings in our domain which should stay fixed.

Training Duration
Our Backbone architecture takes approximately 15 hours, and 400 epochs to converge on a NVIDIA Tesla K80. Some of the more heavyweight models take upwards of 40 hours. In order to conserve our resources and reduce the iteration time, we used a maximum of 150 epochs during our architecture search which amounted to an 6-15 hour period. Most systems had not converged by this point but we found that, for the sample of models that we trained to convergence, the performance ordering at convergence was almost completely unchanged from 150 epochs (see Figures 17 and 5 17 ). All comparisons between AE architectures in Section 5.3 are made with this constraint and the quoted data assimilation figures in these sections should be used for comparative purposes only. We note that this may bias results towards smaller models that train more quickly but, all things considered, we would prefer a bias in this direction. We only trained a single model for each configuration.

Evaluation metric
When evaluating our systems' data assimilation performance we follow [17] in using the following quantity which we refer to as the 'DA MSE': The equivalent quantity for the background state M SE(x b ) is referred to as the 'ref MSE' and is the value of the M SE(·) before DA has taken place. If a DA MSE is lower than the ref MSE, this implies the approach is performing better than the DA baseline system which always predicts the historical mean x b regardless of the observations. The average ref MSE over the 197 test samples is 1.0001 18 . Unless otherwise stated, in all cases in which a single DA MSE is provided, we give the average over the test set.

Hyperparameters
We trained our models with the Adam optimiser with the Pytorch default parameters of β 1 = 0.9 and β 2 = 0.999. We used a fixed learning rate of 0.0002, He et al. initialization [74], and batch size of 16 for most models as this was largest multiple of eight at which the fp32 model, gradients and data could fit in 11GB of available GPU memory. One exception to this was the GDRN model which would only run at batch size 8. Table 3 gives a high-level summary of the results of our architecture search. Note that we also conducted experiments to find the best L and N values for the ResNeXt-L-N and RAB-L CAEs but presentation of these results is deferred to A.

Residual Block
In this experiment we investigated the effect of RB type on our pool of architectures. We found that no single RB was superior for all systems but, for a given architecture, there were large variations in model performance with RB type. For example, vanilla+CBAM RBs were better than NeXt blocks in five of six cases within the ResNeXt/RBD framework. In comparison, CBAMs significantly harmed performance for Tucodec variants. This may be a result of   interference between the coarse-grained attention mechanism of the CBAM blocks and the highly specific attention in the RAB blocks.  Table 5: DA MSE variation with activation function in models trained for 150 epochs. All models used PReLU activations applied channel-wise.

Activation function
In this experiment we investigated the effect of activation function on the best performing systems from the RB experiments. As the Tucodec models were performing well, we investigated the effect of the different activations on all four RBs 19 The results are shown in Table 5.
In IC, Cheng et al. [53], found that PReLU activations were superior to ReLUs while Ma et al. [76] found that GDNs outperformed ReLUs. We believe we are the first to compare GDNs to PReLUs and found that the latter is superior. Crucially, we found that GDNs were very unstable as the function they compute allows for division by zero. In fact, in two of the three cases in Table 5 for which GDNs performed 'best' the models actually produced inf predictions for one of the 197 test set samples. This process occurred more often earlier in training than later, and with the test data than with the training data but as they work well in the Tucodec backbone, we reasoned that they are unstable when the input distribution is unpredictable. Figure 5: The MSE DA with L1 fine-tuning from epoch 300 onwards. The training-set DA MSE decreased but this was not accompanied with a test-set decrease. We investigated L1 fine tuning for all models in Table 6 but just present a representative selection here.

L1 fine tuning
We experimented with the use of L1 fine-tuning late in the training process as recommended in [32] but found that it did not give an appreciable benefit. In fact, as shown in Figure 5, in our experiments it increased the degree of over fitting without providing any generalisation advantage.  Table 6: Summary of the DA MSE and inference speeds of our best performing models after training to convergence.

Architecture
We trained a selection of our best models to convergence and found that the Tucodec-NeXt and Tucodec-vanilla models performed best as shown in Table 6. The two models have very similar DA MSE values but the NeXt model is almost x2.5 faster during inference. As such, we use the Tucodec-NeXt model when making comparisons with reduced space DA in the following section.

Evaluation
In this section, we compare our system against Reduced space VarDA with TSVD as described in [17]. Our system has superior DA performance on the test set as shown in Table 7. This is not just true on average: our system is consistently better (see Figure 6) and space (see B). Moreover, we show in Figure 7 that our method has a DA MSE that is 15% lower than the reduced space approach with M = n even in the limit in which the method becomes Parish et al.'s approach as there is no truncation of V (τ = S). This is surprising: even if our CAE was truly lossless, which it is not, the matrices V and V l contain the same information (albeit the latter stores it more efficiently) 20 . The better performance of our bi-reduced space approach might be explained by the poor conditioning in the mono-reduced space and the resulting numerical errors. This requires further research. Tucodec-NeXt 0.0787 0.0537 Table 7: Comparison of our best Tucodec model with the Arcucci et al. approach [17] which sets σ τ = √ σ 1 = 32.

Model
Our DA MSE is 37% lower than the best Arcucci et al. system. Figure 6: Comparison of TSVD (τ = 32, M = n) and AE data assimilation performance across sequential Fluidity time-steps. Note in Figure a) that the two methods find the same states difficult. In Figure b) we give the difference between the DA MSEs of the two methods for the whole test set. The proposed method performs better (is above the red line) in the vast majority of cases.

Performance-speed tradeoff
The reduced space approach in [17] has an acute performance-speed tradeoff occurring along three axes: 1. The size of the truncation parameter τ . As this increases, the DA performance increases but the speed decreases as shown in Figure 7. 2. The number of observations M . As this increases, the DA performance increases but the speed decreases as shown in Figure 8. This is a stronger effect than 1. Noting the logarithmic scales in this Figure 8b, it is clear that there is only a small range of M for which our method is slower than TSVD. 20 As an aside, note that in this case our method is x43 faster than the traditional approach.  3. The size of the observation variance σ 0 . We do not consider this here but [17] showed that as this parameter increases, performance increases but speed decreases. In all experiments here we used σ 0 = σ l = 0.005.
Our system's evaluation speed is not sensitive to the number of observations 21 , nor the value of σ l . We show the performance-speed tradeoff for a range of models in Figure 9. All timing measurements were carried out on the Intel Xeon E5-2690 v3 (Haswell) 2.60 GHz CPU and averaged over the test set. The clock was started at the beginning of the minimisation routine when all relevant data was already in memory.

Discussion
In this section we briefly discuss a few points raised in the course of this research.

Optimal τ and M
Considering Figure 9, it appears that, for the combinations of τ and M that we investigate here, the most successful pairing in the performance-speed tradeoff is M = 4096 and τ = 791 (i.e. no truncation). It is worth making a few observations on this result: Figure 9: Performance-speed tradeoff for a range of systems.
1. There was no way to know that this combination was the best in advance as it is data-set dependent. It took 60 CPU hours to calculate the DA MSE on the test-set for the range of τ and M displayed in Figure 9. We note that this value could be reduced by a more intelligent search method, but draw the reader's attention to the fact that this is of the same order as the 15 hours required to train a Tucodec-NeXt model to convergence on a GPU.
2. This result is still 30% slower and 20% less accurate for DA than the Tucodec-Next model.

Hardware Accelerators
All timing comparisons were made on the CPU as we did not have a GPU implementation of Arcucci et al.'s routine. Using a K80 GPU with our method provided a speed-up of approximately 40% for our method. This was with a very poor implementation in which the data was transferred from the CPU to the GPU and back again. We expect an optimised implementation of our system on a modern accelerator to achieve a much larger relative speed-up over the figures here than the equivalent optimised version of Arcucci et al.'s method. A full defense of this claim might take another paper but we briefly sketch our our argument for this claim in the following paragraph.
Recall that the bottleneck in our system is a fully convolutional decoder while, in the Arcucci et al. routine, the cost is dominated by large matrix multiplications and vector dot-products. The convolutional kernel parameters are used many times in a forward pass but there is almost no data reuse in the Arcucci et al. case. As such, the latter will be bandwidth-limited but the former may enter the compute-bound domain on some hardware-platforms 22 . Historically, it has proven easier to accelerate compute-bound processes than memory bound ones and there is reason to believe that this will continue, not least because there is a whole industry built-around the design of systems that specifically accelerate convolutional inference workloads. We will not attempt to review the hardware options here but if the 22 Convolution may be memory-bound depending on channel size and dimensions of the feature map.
Graphcore 'IPU', which is set to ship early in 2020, delivers on its marketing promises [77], it would speed up our inference by up to three order of magnitude. There are also lower-cost, lower-power options such as FPGAs [78]. As such, we believe the quoted figures underestimate the latency advantage of our approach.

Other Acceleration Options
During our architecture search, we optimised for DA performance rather than inference latency. Had we been focusing for the latter, there are a number of techniques aside from hardware acceleration that could be used to aid this. Firstly, a thinner decoder could be used as suggested in [6] since only the decoder is evaluated in the online setting. Secondly, the existing network could be quantized [79] or pruned [80,81] or both [?] to provide a substantial speed-increase. Additionally, convolutional acceleration approaches such as Pixel Shuffle [83] or factorised convolutions [84] might be employed to reduce the number of FLOPs in the forward decoder pass. Finally, there are innumerable small architectural changes that could be made in a similar vein to the replacement of vanilla blocks in the original Tucodec model with NeXt blocks. None of the above strategies are available to traditional VarDA approaches. We note that some of these techniques will reduce the performance of our system but, as our approach has a considerable performance cushion over traditional approaches in its current form, this may be acceptable in some settings.

Importance of Architecture
The results in this paper demonstrate the central importance of using good CAE architectures. This field is moving exceptionally fast: our Backbone network, was state-of-the-art for image compression in 2017 [6] but gives a DA MSE that is a) double that of the Tucodec models and b) considerably poorer for DA than the Arcucci et. al. approach with τ = 32 and M = 0.01.
We found that it was non-trivial to extend many architectures to three spatial input dimensions and it required a large amount of manual tuning of the channel sizes so as not to create unreasonably large 4D feature maps (three spatial dimensions and one channel). In particular, our implementation of the GDRN [21] had extreme computational requirements in 3D which, despite its modest number of parameters (see Table 6), took almost three times longer to train than any other network.

Conclusions and Future Work
We have presented a new Bi-reduced space 3D-VarDA formulation and show that, in combination with the Zhou et al. or 'Tucodec' image compression CAE, this method gives superior data assimilation performance in comparison with reduced space VarDA regardless of the parameters used in the latter case. We have demonstrated that our method is also faster in the majority of scenarios. On the theoretical side, we show that our method produces approximately equivalent solutions to the traditional method at lower computational complexity. Unlike the previous approach which is in O(M 2 ) for large M , our method does not penalize the collection of more observation data. We have released our work in a well tested Python module VarDACAE.
There were many extensions to this work which we would have liked to explore further. We feel that the most important of these is the validation of our hypothesis that is possible to create an observation encoder network f o to calculate the latent misfits d l . We would also have liked to apply our approach to 4D-VarDA, validate it on other data sets and investigate alternatives to the L-BFGS minimization routine. A more substantial extension would involve integrating our method with CAE-based ROM approaches to produce a single end-to-end network for reduced space data assimilation and we believe this would be complemented by the use of data assimilation localization techniques [64]. Finally, there is also potential for the use of VAEs within the proposed system to enforce orthogonality in the CAE latent dimension. Figure 10: a) The Residual Dense Block [85], [51] and its extension b) The Grouped Residual Dense Block [21]. Figure 11: The Convolutional Block Attention Module [50]. The channel mask M C (x) and spatial mask M s (x) are applied sequentially. These masks are broadcast to full dimensions (N, C, H, W, D) before their element-wise multiplication with the inputs x. Note that in M C (x), the features from max pooling and average pooling are fed through the same fully connected network one after the other and the results are then added. 'conv n1 7x7x7 p333 s111' represents a convolutional layer with 1 channel, kernel size = (7, 7, 7), padding = (3,3,3) and stride = (1,1,1) and is specific to our implementation (although in some cases we found that kernel size = (3, 3, 3) was necessary to enable efficient training). Our CBAM has just two fully connected layers in M C (x). Figure 12: The Residual Attention Block proposed by [51] and utilised by [20]. Note that unlike CBAMs, the trunk of RABs (yellow background) are not the identity mapping. Table 1 became available to search. We evaluated the grid search of options within this space and find that 27 layers of Figure 13: The GRDN block [?] in a series of GRDBs (see in Figure 10) are used with a CBAM module. Figure 14: a) A single ResNeXt layer, repeated from Figure ?? for clarity [23]. The ResNeXt cardinality describes the number of RBs in each layer, b) three stacked ResNeXt layers with an extra residual connection, and c) a ResNeXt layer with RDBs instead of simple RBs. Each 'RDB3' has 3 RBs. Attenuation coefficient We found that it was difficult to train ResNeXt variants with large cardinality but, as the backbone trained easily, it was clear that the new residual blocks were interfering with the backbone's ability to downsample the inputs. Therefore we introduced a residual attenuation coefficient α at the exit to the block shown in Figure 15 such that the computed function is: g(x) = x + α f (x) (71) α was initialised to 0.05 at the start of training and then updated with the other network parameters. This down-weights the ResNext block's importance initially so that the backbone has time to learn a good compression.
ResNeXt width and cardinality In the first experiment, we investigated the effect of changing cardinality and number of layers within residual component of our ResNeXt variant. We used NeXt RBs for these experiments. The results are in Figure 16. The best system was ResNeXt3-27-4-NeXt but there was also an interesting group of models with three layers on the left of Figure 16. We decided to preserve this diversity and take three models forward to the next stage of experiments. These were ResNeXt3-27-4-NeXt, ResNeXt3-3-8-NeXt and ResNeXt3-27-1-NeXt 23 .  In our second experiment we investigated the effect changing the number of consecutive RAB blocks in the backbone network. The results are in Table 8. The the best of these systems with four RABs, gives a 17% improvement relative to the Backbone but is poor in comparison with the best Tucodec model.  Table 9: The field-jitter augmentation strengths we investigated. We added Gaussian noise with standard deviation of 'Jitter Amplitude' at 'Jitter Frequency' of the total locations in the state.

RAB-
To quantify what, if any, effect our augmentation technique was having, we retrained the Tucodec-NeXt model with a range of augmentation strengths as shown in Figure 17. We did not observe a large difference between the methods so choose the strongest augmentation that did not harm performance (augmentation strength 2 in Table 9) when training our models to convergence.  Table 9. These graphs have been smoothed with an exponential moving average with α = 0.4 as the spikes in the training curves made this diagram too noisy to be illustrative. A non-smoothed version is given in Appendix ??.

B Further Comparisons
This Appendix contains two graphs that would have been repetitious in the full text but provide useful context to the comparison between reduced space VarDA and bi-reduced space VarDA.   Figure 6 with τ = 791 (i.e. no truncation) instead of τ = 32 and M = n as before. Although the performance is more similar in this case, our method still performs better on average. We also note that, in order to achieve this level performance, the reduced space method takes 2.5s, or x43 longer than our approach.