# Are discrete units necessary
for Spoken Language Modeling?

Tu Anh Nguyen, Benoit Sagot, Emmanuel Dupoux Tu Anh Nguyen is with Meta and
Inria, France, e-mail<EMAIL_ADDRESS>Benoit Sagot is with Inria,
France, e-mail<EMAIL_ADDRESS>Emmanuel Dupoux is with Meta and EHESS,
ENS-PSL, CNRS, Inria, France, e-mail<EMAIL_ADDRESS>
###### Abstract

Recent work in spoken language modeling shows the possibility of learning a
language unsupervisedly from raw audio without any text labels. The approach
relies first on transforming the audio into a sequence of discrete units (or
pseudo-text) and then training a language model directly on such pseudo-text.
Is such a discrete bottleneck necessary, potentially introducing irreversible
errors in the encoding of the speech signal, or could we learn a language
model without discrete units at all? In this work, we study the role of
discrete versus continuous representations in spoken language modeling. We
show that discretization is indeed essential for good results in spoken
language modeling. We show that discretization removes linguistically
irrelevant information from the continuous features, helping to improve
language modeling performances. On the basis of this study, we train a
language model on the discrete units of the HuBERT features, reaching new
state-of-the-art results in the lexical, syntactic and semantic metrics of the
Zero Resource Speech Challenge 2021 (Track 1 - Speech Only).

###### Index Terms:

Spoken Language Modeling, Discrete Units, HuBERT

##  I Introduction

Pre-training language models on large-scale text data have achieved tremendous
success in natural language understanding and have become a standard in
Natural Language Processing (NLP) [1, 2, 3, 4, 5]. Recently, [3] showed that
very large language models are actually few-shot learners, and manage to
perform well even in zero-shot settings.

Large-scale self-supervised pre-training for speech data has also become more
and more popular as a method to boost the performance of Automatic Speech
Recognition (ASR) [6, 7, 8]. However, these models mostly rely on fine-tuning,
which requires more training and text labels, to either improve the model or
evaluate the learned representations of the speech. Lately, [9] introduces a
new unsupervised task: Spoken language modeling, the learning of a language
unsupervisedly from raw audio without any text labels, along with a suite of 4
zero-shot metrics probing for the quality of the learned models at different
linguistic levels: phonetic, lexical, syntactic, semantic. The metrics are
evaluated using the representations extracted from the model (phonetic,
semantic) or pseudo-probability scores given by the model (lexical,
syntactic). Their proposed baseline approach relies on transforming the audio
into a sequence of frame-by-frame discrete units (or pseudo-text) and training
a language model on the pseudo-text. The trained models displayed better-than-
chance performances on nearly all the evaluation metrics of the challenge [9,
10]. However, this paradigm creates a discrete bottleneck between a speech
encoder and a language model which could be a potential source of error, and
in addition requires multiple training phases (learning an acoustic
representation, clustering it, and learning a language model). Is such a
discrete bottleneck necessary?

One way in which discrete units could help language modeling stems from the
fact that in contrast to text, audio data contains a lot more details, some of
which are linguistically relevant (intonation, rhythm, non verbal
vocalization), others not so (background noise, reverberation, speaker
identity, etc). To the extent that discretization effectively removes
linguistically irrelevant information from the continuous features [11], it
could indeed help language modeling. Of course, this potential gain could be
counterbalanced by the fact that discretization could also make errors and
remove useful information.

Figure 1: Overview of the trained BERT models. The BERT model takes as input
either the continuous features extracted from CPC or the sequences of frame-
by-frame discretized units obtained from k-means, and tries to predict either
continuous target features (with L1, L2 or NCE loss) or discrete target units
(with NLL loss).

In this work, we analyse the importance of discretization in spoken language
modeling. We employ a pre-trained acoustic model to obtain either continuous
or discretized features from audio data. We then train BERT language models
with a Masked Language Modeling (MLM) objective on both discrete and
continuous features used either as inputs or as targets and evaluate the
resulting systems on zero-shot spoken language modeling metrics. We also
evaluate HuBERT [7], a single model trained from raw waveform with discrete
targets, on these metrics and compare the results with our best models.

Our contributions can be listed as follows:

  * •

We show experimentally that discretization is beneficial for spoken language
modeling, but we can get rid of discrete bottlenecks by using low-level
continuous inputs so long as we still use discrete targets.

  * •

We show that discretization disentangles linguistic information from non-
linguistic signals, forcing the transformer to focus on linguistic ones.

  * •

We show that a self-supervised model trained with a MLM objective on discrete
targets like HuBERT achieves very good results on spoken language modeling
metrics, showing that it can learn not only acoustic but also high-level
linguistic information.

##  II Related Work

#### Discretization in Self-Supervised Approaches

Self-supervised models for learning speech representation have become more and
more popular as an effective pre-training method for downstream Automatic
Speech Recognition (ASR) task, notably wav2vec2.0 [6] and HuBERT [7]. Both
models comprise a feature extractor (CNN Encoder) followed by a feature
encoder (Transformer Encoder), and are trained with a MLM objective like BERT.
However, wav2vec2.0 discretizes the latent features obtained by the CNN
Encoder and uses them as the target for the Transformer Encoder using a
contrastive loss against negative samples in the sentence. On the other hand,
HuBERT discretizes fixed features obtained from a teacher model and uses these
fixed discrete units as the target for the Transformer Encoder using a cross-
entropy loss. Finally, our work is mostly similar to [12], where they compare
BERT models training on discrete units obtained from vq-wav2vec [13] and
continuous features obtained from wav2vec [14] on the ASR task. They found
that training BERT model on discrete vq-wav2vec units is more effective for
ASR.

#### Spoken Language Modeling

Following the huge success of language models on text data [4, 2, 3], the Zero
Resource Speech Challenge 2021 [9, 10] opens up new possibilities for learning
high-level language properties from raw audio without any text labels. They
introduced 4 zero-shot evaluation metrics at different linguistic levels
(phonetic, lexical, syntactic, semantic), along with composite baseline
systems consisting of an acoustic discretization module (Contrastive
Predictive Coding, or CPC+k-means) followed by a language model (BERT or LSTM)
on the discretized units. The CPC model takes the raw audio as input and
produces phonetic representations at a lower frame rate of 100Hz, helping the
language model to learn high-level information from the raw audio. In the same
spirit, [15] introduced Generative Spoken Language Modeling (GSLM), the task
of learning and generating spoken language from raw audio only. They provided
baseline systems consisting of a discrete speech encoder (CPC, wav2vec 2.0,
HuBERT), a generative language model (GPT-like model), and a speech decoder
(Tacotron-2, [16]). The models are evaluated on spoken language modeling
metrics [9], ASR-based generation metrics [15] as well as human evaluation
metrics.

##  III Experimental Setup

In this section, we first present the evaluation metrics as well as the
dataset used to train and evaluate the models. We then explain our models and
the inference methods for model evaluation.

###  III-A Evaluation Metrics

We evaluate our models with the ZeroSpeech 2021 Benchmark Metrics [9],
consisting of 4 zero-shot tests probing for the quality of spoken language
models at four linguistic levels: phonetic (Libri-light ABX metrics), lexical
(sWUGGY spot-the-word metrics), syntactic (sBLIMP acceptability metrics) and
semantic (sSIMI similarity metrics).

#### Libri-light ABX metrics

Given a pair of similar triphones (e.g., ‘aba’-‘apa’) spoken by a same speaker
and an intervening sound (either ‘aba’ or ‘apa’), the model has to tell which
sound has a closer representation to the intervening sound. The ABX metrics is
reported as the error rate that the model fails to choose the correct
triphone.

#### sWUGGY spot-the-word metrics

Given a pair of a word and a similar non-word (e.g., ‘brick’-‘blick’), the
model has to tell which is the word based on their probability. The spot-the-
word metrics is reported as the accuracy that the model assigns a higher
probability to the word.

#### sBLIMP acceptability metrics

Given a linguistic minimal sentence pair of matched grammatical and
ungrammatical sentences (e.g., ‘he loves it’-‘he love it’), the model has to
tell which is the grammatical sentence. The acceptability metrics is reported
as the accuracy that the model assigns a higher probability to the grammatical
sentence.

#### sSIMI similarity metrics

Given a pair of words (e.g., ‘happy’-‘joyful’), the model has to compute a
similarity score based on their representations. The similarity metrics is
reported as the Pearson correlation coefficient (PCC) between model scores and
human judgements. In this work, the sSIMI scores are weighted across different
subsets according to their sizes and averaged across LibriSpeech and synthetic
subsets to make it more accurate and consistent. We reported it as wSIMI.

###  III-B Datasets

#### Training Dataset

We train our models on LibriSpeech [17], an English corpus containing 1000
hours of read speech based on public domain audio books. The models are
validated on LibriSpeech dev-clean and dev-other subsets, comprising 10 hours
of speech in total.

#### Metrics Datasets

The metrics datasets are either extracted sounds from LibriSpeech (ABX, sSIMI)
or synthesised using Google API111https://cloud.google.com/text-to-speech
(sWUGGY, sBLIMP, sSIMI). The datasets containing words or sentences were
filtered to only contain the LibriSpeech vocabulary (except sWUGGY non-words),
and are split into dev and test sets. The dev sets have been made publicly
available at the ZeroSpeech 2021 Challenge website
222https://zerospeech.com/2021/instructions.html#evaluation-dataset.

###  III-C Models

#### ZeroSpeech 2021 Baseline

The ZeroSpeech 2021 Baseline System [9] is a composite of three components: an
acoustic model (CPC, [18, 19]), a clustering module (k-means) and a language
model (BERT, [4]). The CPC model is first trained to obtain good phonetic
representations of the speech, which are then discretized into sequences of
units with the k-means model. The BERT model is finally trained on these
discrete units to better learn linguistic information.

As we only focus on the language modeling system in this work, we shall use
the best CPC model in the ZeroSpeech 2021 Baseline System, which comprises a
5-layer 1D-CNN Encoder followed by a 4-layer LSTM autoregressive model. The
features are extracted from the 2nd layer (unless otherwise specified) of the
LSTM model, with a rate of 100Hz, and are either discretized with a 50-unit
k-means model (discrete) or left unchanged (continuous).

#### BERT with discrete and continuous features

We modify the BERT model so that it is able to take as input either discrete
units obtained from k-means or continuous features extracted from CPC, in
which case the masking is done by replacing the features with a masked
embedding vector. We also allow the model to predict either discrete target
units or continuous target features, with multiple choices of an appropriate
objective for each case. When predicting discrete targets, we use a cross-
entropy objective (Negative Log-Likelihood, or NLL loss) but with two slightly
different implementations. We could simply employ a linear classification head
at the output of the BERT model as usual (which we denote by linear NLL, or
NLL-l) or force the BERT output features to be similar to the embedding
vectors of the target units as for HuBERT (cf. equation (3) from [7], we
denote this by embedding NLL, or NLL-e). In the case of continuous targets, it
can be a reconstruction objective (L1 loss or L2 loss) or a contrastive
objective (Noise Contrastive Estimation, or NCE loss). In the latter case, the
predicted features are contrasted with 100 negative features sampled from the
same phrase (similar to continuousBERT, [12]).

We use a BERT base model, which comprises a 12-layer Transformer Encoder. Our
implementation is based on the wav2vec2.0 [6] Transformer Encoder
333https://github.com/pytorch/fairseq/tree/main/examples/wav2vec using fairseq
[20]. Each input sequence contains the features of a full audio file, and we
consider at most 15.6 seconds of audio per file. We trained all models for
250k update steps on 32 GPUs, with a batch size of 175s per GPU. The learning
rate was warmed up to a peak value of $1\times 10^{-5}$ after 32k steps. For
the masking, we masked $M$ consecutive tokens for each span, where
$M\sim\mathcal{N}(10,10)$, with a total masking coverage of roughly half of
the input tokens (spans may overlap).

###  III-D Model Inference for Evaluation

#### ABX Distance

For the ABX metrics, we extract frame-by-frame representation features for
each audio file. Then, the ABX distance between two files is computed as the
average angular distance of the representations along the realigned Dynamic
Time Wrapping path. Given two audio files $x$ and $y$ with two sequences of
representation $\textbf{r}^{x}=r^{x}_{1},\dots,r^{x}_{T}$ and
$\textbf{r}^{y}=r^{y}_{1},\dots,r^{y}_{S}$ respectively, the ABX distance
between $x$ and $y$ is computed as follows:

$\begin{split}&d_{ABX}(x,y)=\frac{1}{\lvert\text{path}_{\text{DTW}}(\textbf{r}^{x},\textbf{r}^{y})\rvert}\sum_{(i,j)\in\text{path}_{\text{DTW}}(\textbf{r}^{x},\textbf{r}^{y})}sim(r^{x}_{i},r^{y}_{j}),\end{split}$
(1)

where $sim(r^{x}_{i},r^{y}_{j})$ is the angular distance (in radian) between
the embeddings $r^{x}_{i}$ and $r^{y}_{j}$.

We note that in this paper the ABX metrics are mainly used to evaluate the
input and target features of the BERT model, and therefore the ABX distances
are mostly performed on the CPC features without using the BERT model.

#### Probability Estimation

For sWUGGY and sBLIMP metrics, we compute for each audio file a model-based
pseudo log-probability (m-PLP) of the trained BERT model. Given an audio file
$x$ with the input and target features for the BERT model $x_{1}...x_{T}$ and
$\hat{x}_{1}...\hat{x}_{T}$ respectively, the m-PLP is computed as follows:

$\begin{split}&\text{m-PLP}(x)=\quad\sum_{\begin{subarray}{c}j=0\\\ i=j\Delta
t\end{subarray}}^{\lfloor(T-M)/\Delta
t\rfloor}\sum_{m=1}^{M}PLP(\hat{x}_{i+m}|\overline{x_{i+1}..x_{i+M}}),\end{split}$
(2)

where $M$ is a chosen size of a sliding window, $\Delta t$ is a chosen step of
the sliding window and $PLP(\hat{x}_{i+m}|\overline{x_{i+1}..x_{i+M}})$ is a
pseudo log-probability of the target $\hat{x}_{i+m}$ given by the BERT model
with $M$-span masked inputs $x_{1}..x_{i}m..mx_{i+M+1}..x_{T}$ ($m$ represents
a masked feature).

For models with NLL or NCE loss,
$PLP(\hat{x}_{i}|\overline{x_{i+1}..x_{i+M}})$ is computed as the log value of
the probability given by the softmax layer of the BERT model (in the NLL case,
the probability is computed over all tokens, while in the NCE case it is
computed over all sampled negative examples). For models with L1 or L2 loss,
we compute $PLP(\hat{x}_{i}|\overline{x_{i+1}..x_{i+M}})$ as the negative
reconstruction loss of the predicted feature and the target feature
$\hat{x}_{i}$. The negativity ensures that a correct target has a higher
m-PLP.

The m-PLP extends the span-masked pseudo probability (span-PP) [9] to BERT
models with continuous targets. It is derived from the pseudo-loglikelihood
score (PLL) for MLMs [21], which was shown to be an effective sentence scoring
method for BERT models in many scenarios [22].

The choice of $M$ and $\Delta t$ is determined for each model using the dev
sets, and is given in Table VI. In our experiments, we always consider $\Delta
t=5$ and vary $M$ in $\\{15,25,35,45,55\\}$. For models trained on HuBERT
features (section IV-C), we vary $M$ in $\\{5,10,15,20,25\\}$ as the frame
rate is 50Hz instead of 100Hz as for CPC.

#### Similarity Score

For the sSIMI metrics, we extract a fixed-length representation for each audio
file by applying a pooling function (mean, max, min) over hidden features from
one layer of the Transformer Encoder. The similarity score of two audio files
is computed as the cosine similarity between the two corresponding
representations. The choice of the hidden layer and the pooling function is
determined for each model using the dev sets and is given in Table VI.

##  IV Results

###  IV-A Discrete bottleneck seems to be essential for spoken language
modeling

Table I reports the performances of our BERT models, trained with either
continuous or discrete CPC features of the LibriSpeech 960h dataset, on
lexical (sWUGGY), syntactic (sBLIMP) and semantic (wSIMI) metrics.

id input target loss sWUGGY↑ sBLIMP↑ wSIMI↑ discrete input, discrete target 1
disc. disc. NLL-l 79.28 59.71 6.32 2 disc. disc. NLL-e 80.02 59.86 7.87
continuous input, discrete target 3 cont. disc. NLL-l 60.36 53.23 8.39 4 cont.
disc. NLL-e 60.20 52.78 9.49 continuous input, continuous target 5 cont. cont.
NCE 56.84 52.62 9.16 6 cont. cont. L1 59.23 53.12 7.85 7 cont. cont. L2 60.56
53.33 6.55 discrete input, continuous target 8 disc. cont. NCE 65.69 57.24
9.33 9 disc. cont. L1 73.93 56.02 10.69 10 disc. cont. L2 74.22 55.75 5.97

TABLE I:  Performances on the dev sets of sWUGGY, sBLIMP, wSIMI metrics of
BERT models using either continuous ZeroSpeech CPC features (layer 2 of the
LSTM module of CPC-big) or discretized features (with a 50-unit k-means model)
as inputs and targets. Best scores in each category are in bold, best scores
overall are underlined.

We first examine how the continuity of the input and target features affects
the quality of the BERT model on the evaluation metrics. By comparing the best
scores in each case, we see that having discrete inputs helps the model learn
better lexical and syntactic information, whereas models with continuous
inputs do have better than chance performance on the lexical task. We observe
that the best models on the language model tasks are obtained with discrete
inputs and discrete targets, which is the classic configuration of BERT.
Predicting continuous targets from discrete inputs, where the model acts as an
autoencoder decoder, is also beneficial and nearly catches up with the best
models. It is interesting, still, to note that it is possible to acquire some
language information without any discretization. The wSIMI scores are still
quite low, but we see in general that having continuous information does help.

###  IV-B Is continuous input always bad?

We observe during our training experiments that the masked prediction
objective is too easy for some models with continuous inputs and could quickly
lead to overfitting. This could be explained by the fact that the input and
target features are extracted from the same layer of the LSTM autoregressive
module of CPC. As a consequence, we try using the input features from
different layers of the LSTM module, while maintaining the same target layer.
We keep using the NLL-e loss for discrete targets while using NCE and L1 loss
for continuous targets. The results are reported in Table III.

ABX within↓ ABX across↓ clean other clean other layer 0 cont. 11.50 14.09
18.53 24.70 disc. 21.46 24.21 30.77 34.91 layer 2 cont. 3.41 4.84 4.20 7.65
disc. 6.38 10.22 8.22 14.86 layer 4 cont. 9.49 11.95 10.01 15.70 disc. 19.81
21.64 24.39 28.04

TABLE II: Within and Across Speaker ABX error (lower is better) on Libri-light
dev-clean and -other for continuous and discretized features of different
layers of the LSTM autoregressive module of CPC-big model. Layer 0 means the
output of the CNN Encoder module.

We observe that using continuous input features from a different layer does
reduce overfitting during training, which significantly improves the
performances of the models on LM metrics, especially for sWUGGY scores.
Interestingly, we note that using continuous input features from a lower LSTM
layer (layer 0, where the ABX errors are high, cf. Table II) to predict target
features from a higher LSTM layer (layer 2) is more beneficial to the model
than using high quality continuous input features from the same or higher
layer as the target features (layer 2, layer 4). This is not the case,
however, for discrete input models, where the model benefits from good quality
input units.

id input target sWUGGY↑ sBLIMP↑ wSIMI↑ discrete input, discrete target, NLL-e
loss 2 layer 2 layer 2 80.02 59.86 7.87 11 layer 0 layer 2 64.91 52.45 7.48 12
layer 4 layer 2 70.68 55.06 8.61 continuous input, discrete target, NLL-e loss
4 layer 2 layer 2 60.20 52.78 9.49 13 layer 0 layer 2 77.19 55.30 7.25 14
layer 4 layer 2 67.41 54.13 8.06 continuous input, continuous target, NCE loss
5 layer 2 layer 2 56.84 52.62 9.16 15 layer 0 layer 2 65.53 55.20 6.17 16
layer 4 layer 2 59.81 52.93 8.32 continuous input, continuous target, L1 loss
6 layer 2 layer 2 59.23 53.12 7.85 17 layer 0 layer 2 67.83 53.59 7.25 18
layer 4 layer 2 63.68 53.26 6.90 discrete input, continuous target, NCE loss 8
layer 2 layer 2 65.69 57.24 9.33 19 layer 0 layer 2 58.55 52.31 8.07 20 layer
4 layer 2 58.61 54.48 7.72 discrete input, continuous target, L1 loss 9 layer
2 layer 2 73.93 56.02 10.69 21 layer 0 layer 2 62.98 53.47 5.20 22 layer 4
layer 2 65.92 53.94 7.01

TABLE III:  Performances on the dev sets of sWUGGY, sBLIMP, wSIMI metrics of
BERT models using the input features from different layers of the LSTM module
of CPC-big model. Layer 0 means the output of the CNN Encoder module. Best
scores in each category are in bold, best scores overall are underlined.

Overall, we observe that discrete-discrete model (with the same input and
target units) yields the best performance when the quality of discrete units
is good. Continuous-discrete is also a great choice when using low-level input
features. When there are no discrete units at all, the LM performances are
still limited, even if using a NCE loss could help a bit with syntactic and
semantic metrics.

###  IV-C Varying the number of discrete units

Here, we address the question as to why discrete units are better than
continuous ones. One hypothesis is that discrete units manage to remove
linguistically irrelevant information and force the transformer to focus on
linguistic ones. To test this, we run a speaker discrimination probe on the
discrete units and continuous features. In addition, we run a new experiment
varying the number of discrete units from 20 to 2000. Hypothetically, when the
number of units is too small (eg, smaller than the number of phonemes), the
resulting phonetic confusions should degrade the learning of higher linguistic
representations. Conversely, when the number of units is too large, the
quantization step would start to leak other-than-phonetic information into the
representation, hence making it closer to the continuous representations.

unit quality language modeling on units model n units spk prb↑ ABX↓ sWUGGY↑
sBLIMP↑ sSIMI↑ CPC 20 30.40 12.66 71.71 58.97 4.72 50 34.00 9.89 80.02 59.86
7.87 100 49.20 9.56 80.47 59.47 6.09 200 56.00 9.72 79.90 58.90 4.45 500 64.00
10.72 79.66 59.72 6.37 1000 61.60 11.99 79.86 58.46 6.78 2000 67.60 14.24
78.46 58.25 5.94 cont. 98.00 5.02 - HuBERT 20 24.40 14.04 62.89 57.06 7.80 50
38.00 9.19 76.79 61.12 8.61 100 48.00 8.34 81.09 62.47 5.19 200 61.60 7.57
81.54 62.78 7.03 500 68.40 7.73 83.06 62.89 9.73 1000 74.40 9.04 82.58 61.55
8.64 2000 73.20 11.00 81.61 62.85 10.66 cont. 99.60 4.23 - Forced Phones 40
10.00 0.00 92.19 63.72 6.23

TABLE IV: Discrete Unit quality (Speaker probing and ABX) and Performance of
the BERT models trained on Discrete Units on the dev sets of LM scores
(sWUGGY, sBLIMP, wSIMI) for different numbers of clusters on CPC and HuBERT
features. The ABX is averaged on dev-clean and dev-other within and across
subsets.

To support our hypothesis, we run kmeans on the continuous features of both
CPC and HuBERT models, and vary k to be 20, 50, 100, 200, 500, 1000, and 2000,
after which we train a discrete-discrete BERT model. For the CPC features, we
take the layer 2 features of the CPC-big model as usual. For the HuBERT
features, we train our own HuBERT base model as described in Section IV-D, we
then take the features from layer 12 of the Transformer Encoder after the 2nd
iteration, which have the best ABX (cf. Table VII). Following [23], we train a
speaker classifier in the following way: We randomly split LibriSpeech dev-
clean utterances into train/valid/test (80%/10%/10%) sets and train a two-
layer Transformer classifier on the sequences of discrete units or continuous
features of the utterances. The classification head is performed on the first
token (bos, or begin-of-sentence) of the transformer outputs. For this speaker
probing task, there are 40 classes (speakers). The models are trained for 20
epochs and are validated on the valid set. We finally report the test
accuracy. For reference, we also include the forced phonemes units (frame-by-
frame phonemes). As the forced phonemes contain a silence, there are 40 units
in total.

Figure 2: ABX of Discrete units and Error rate on the dev sets of LM scores
(sWUGGY, sBLIMP) for different numbers of clusters for CPC and HuBERT
features. The ABX is averaged on dev-clean and dev-other within and across
subsets.

The results are reported in Table IV and illustrated in Figure 2. As expected
speaker classification accuracy increases with the number of clusters, and the
continuous features yield the best classification. We can observe a U-shaped
curve in performance across the different language metrics as a function of
the number of units. Interestingly, the optimum number of units seems to be
different across the model features (CPC, HuBERT) and linguistic levels.
HuBERT features are better than CPC features in most cases, and seem to
benefit from more clusters than CPC features. In general, we see that the
language model scores seem to decrease slowly compared to ABX as the number of
clusters becomes bigger. It is also interesting to note that the language
model scores become steadily good as soon as the number of clusters is higher
than the number of phonemes (40 units). This could be seen in Figure 3, where
we observe a limited unit-phoneme correspondence when having only 20 discrete
units; but as soon as the number of clusters reaches 50, we see a clear
correspondence between the units and the phonemes, although several ”hard”
phonemes are still dispersed and don’t correspond to a single unit (e.g. ch,
oy, th, uh); when there are 500 clusters, there are more units representing a
single phoneme, and most ”hard” phonemes are now assigned by certain units.

These results support the hypothesis that the superiority of the discrete
units is due to the fact that they block the propagation and amplification of
non-linguistic signals that may be present (even if attenuated) in continuous
representations.

###  IV-D Comparison with state-of-the-art systems

ABX (target features)↓ Systems within across sWUGGY↑ sBLIMP↑ wSIMI↑ id input
target loss clean other clean other ZeroSpeech 2021 Best Baseline System [9]
CPC-layer2+km50 CPC-layer2+km50 NLL-l 6.71 10.62 8.41 15.06 75.51 56.16 2.05
ZeroSpeech 2021 Text Topline Systems [9] Forced phones Forced phones NLL-l
0.00 0.00 0.00 0.00 91.88 63.16 4.44 Phones Phones NLL-l - - - - 97.67 66.91
12.80 BERT Models on CPC-big Features 2 CPC-layer2+km50 CPC-layer2+km50 NLL-e
6.71 10.62 8.41 15.06 80.29 59.93 6.56 13 CPC-layer0 CPC-layer2+km50 NLL-e
6.71 10.62 8.41 15.06 77.22 55.62 6.61 17 CPC-layer0 CPC-layer2 L1 3.28 4.81
4.31 7.92 68.37 53.95 5.68 9 CPC-layer2+km50 CPC-layer2 L1 3.28 4.81 4.31 7.92
74.46 55.38 6.17 HuBERT Base Models 23 waveform MFCC+km100 NLL-e 20.22 24.97
33.42 40.45 62.74 54.11 5.58 24 waveform H-iter1-layer6+km500 NLL-e 6.29 7.51
8.76 12.82 79.13 58.89 5.45 25 waveform H-iter2-layer12+km500 NLL-e 5.87 7.15
6.96 10.73 80.19 59.29 5.87 BERT Models on HuBERT Discrete Units 26
H-iter2-layer12+km500 H-iter2-layer12+km500 NLL-e 5.87 7.15 6.96 10.73 83.29
61.93 9.73

TABLE V:  Comparison on the test sets of the 4 ZeroSpeech 2021 metrics of our
BERT models trained on continuous or discrete CPC features, BERT model trained
on HuBERT discrete units and HuBERT Base models with ZeroSpeech 2021 Baseline
and Topline Systems. For each continuous/discrete combination, we choose the
best performing model on the dev set as reported in Table III. We trained the
HuBERT model for 3 iterations. The targets used to train the 3 iterations are
discretized MFCC features (100 units), discretized features from Transformer’s
layer6 of 1st iteration (500 units) and discretized features from
Transformer’s layer12 of 2nd iteration (500 units) respectively. All models
were trained on the LibriSpeech 960h dataset. For the ABX metrics, we report
the scores on the target features used to train the model. Best scores in each
category are in bold, best scores overall are underlined.

We evaluate the HuBERT model on the zero-shots metrics and compare the results
with our trained BERT models. The HuBERT model is trained iteratively, using
clustering units from features of previous iteration as the teacher. We
trained a HuBERT base model, which comprises a 7-layer CNN Encoder followed by
a 12-layer Transformer Encoder, on the Librispeech 960h dataset for 3
iterations. The teachers for each iteration are MFCC features (100 units),
Transformer’s layer6 of 1st iteration (500 units) and Transformer’s layer12 of
2nd iteration (500 units) respectively. Architecturally, the Transformer
Encoder of the HuBERT model is very similar to our model 13 (continuous input
layer 0, discrete target layer 2, NLL-e loss) where they both take as input
the continuous features of the CNN Encoder and predict discrete targets
obtained from features of a higher level with a NLL-eloss.

Overall performances on the ZeroSpeech 2021 test sets are reported in Table V.
For each of discrete/continuous combinations, we choose the best performing
model on the dev set as reported in Table III. We also include the discrete-
discrete model trained on HuBERT Discrete Units (500 units), which was
reported to have the best LM scores in section IV-C. We first observe a huge
improvement of model 2 compared with the baseline system, even if they both
use the same units for the BERT model. This improvement greatly comes from the
reimplementation of the BERT model, which uses the wav2vec2 Transformer
Encoder model444One main difference between the two Transformer models is that
wav2vec2 uses a Convolutional Positional Embedding instead of the standard
Sinusoidal Positional Embedding. However, we did not study the effect of this
difference in this paper.. Changing the NLL-l loss to NLL-e loss also improves
a little bit (cf. Table I).

It seems that using good quality discrete units as targets is very beneficial
for the language models, achieving better scores than using continuous targets
in all the metrics. The HuBERT model performs surprisingly well, approaching
our best model on the language model tasks. This means that the Transformer
Encoder of HuBERT acts as a language model as well. We see that as soon as the
discrete targets have better quality, the HuBERT model manages to have better
results on spoken language modeling metrics. We see that the discrete-discrete
model on HuBERT Discrete Units (model 26) further improves the scores on all
the metrics, confirming again our finding that it’s better to train a
discrete-discrete model when we have good quality units.

Comparing the results with the ZeroSpeech 2021 Systems, we observe that our
models are closing the gap between spoken and text-based language models.

##  V Conclusion

This work analyses the importance of discretization in spoken language
modeling. We experimentally show that discretization is essential for spoken
language modeling, although high-quality discrete units are required to obtain
good performances. We also show the possibility of learning high-level
language properties of a self-supervised speech representation learning model
like HuBERT. Finally, we obtain state-of-the-art results on 3 out of 4 metrics
of the Zero Resource Speech Challenge 2021 (Track 1 - Speech Only), bridging
the gap between speech and text-based systems. Note though that because HuBERT
requires a teacher that learns a discrete representation, the overall training
of HuBERT is not end-to-end, because the training of the teacher is not (in
fact, requires several iterations). Further work is needed to simplify this
kind of training loop to learn language directly from speech inputs. Further
work is also needed to assess whether the present results can generalize to
other languages and datasets.

## Acknowledgments

In this work, ED in his EHESS role was supported by the Agence Nationale pour
la Recherche (ANR-17-EURE-0017 Frontcog, ANR-10-IDEX-0001-02 PSL*,
ANR-19-P3IA-0001 PRAIRIE 3IA Institute), a grant from CIFAR (Learning in
Machines and Brains) and HPC resources from GENCI–IDRIS (Grant
2021-AD011011691R1).

We would like to thank the reviewers for their thoughtful comments on the
paper. We also thank Jade Copet, Evgeny Kharitonov, Morgane Riviere, Paden
Tomasello, Wei-Ning Hsu, Yossef Mordechay Adi, Abdelrahman Mohamed, Maureen de
Seyssel, Marvin Lavechin, Robin Algayres, Xuan-Nga Cao, Nicolas Hamilakis,
Hadrien Titeux, Gwendal Virlet, Marianne Metais for helpful discussions on the
paper.

## References

  * [1] A. Radford, K. Narasimhan _et al._ , “Improving language understanding by generative pre-training,” in _arxiv_ , 2018.
  * [2] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever _et al._ , “Language models are unsupervised multitask learners,” _OpenAI blog_ , vol. 1, no. 8, p. 9, 2019.
  * [3] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei, “Language models are few-shot learners,” in _Advances in Neural Information Processing Systems_ , H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, Eds., vol. 33. Curran Associates, Inc., 2020, pp. 1877–1901. [Online]. Available: https://proceedings.neurips.cc/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf
  * [4] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “BERT: Pre-training of deep bidirectional transformers for language understanding,” in _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_. Minneapolis, Minnesota: Association for Computational Linguistics, Jun. 2019, pp. 4171–4186. [Online]. Available: https://aclanthology.org/N19-1423
  * [5] Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoyanov, “RoBERTa: A robustly optimized BERT pretraining approach,” _CoRR_ , vol. abs/1907.11692, 2019. [Online]. Available: http://arxiv.org/abs/1907.11692
  * [6] A. Baevski, Y. Zhou, A. Mohamed, and M. Auli, “wav2vec 2.0: A framework for self-supervised learning of speech representations,” in _Advances in Neural Information Processing Systems_ , H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, Eds., vol. 33. Curran Associates, Inc., 2020, pp. 12 449–12 460. [Online]. Available: https://proceedings.neurips.cc/paper/2020/file/92d1e1eb1cd6f9fba3227870bb6d7f07-Paper.pdf
  * [7] W.-N. Hsu, B. Bolte, Y.-H. H. Tsai, K. Lakhotia, R. Salakhutdinov, and A. Mohamed, “HuBERT: Self-supervised speech representation learning by masked prediction of hidden units,” _IEEE/ACM Transactions on Audio, Speech, and Language Processing_ , vol. 29, pp. 3451–3460, 2021. [Online]. Available: https://ieeexplore.ieee.org/document/9585401
  * [8] Y. Chung, Y. Zhang, W. Han, C. Chiu, J. Qin, R. Pang, and Y. Wu, “W2v-BERT: Combining contrastive learning and masked language modeling for self-supervised speech pre-training,” _CoRR_ , vol. abs/2108.06209, 2021\. [Online]. Available: https://arxiv.org/abs/2108.06209
  * [9] T. A. Nguyen, M. de Seyssel, P. Rozé, M. Rivière, E. Kharitonov, A. Baevski, E. Dunbar, and E. Dupoux, “The zero resource speech benchmark 2021: Metrics and baselines for unsupervised spoken language modeling,” in _Neural Information Processing Systems Workshop on Self-Supervised Learning for Speech and Audio Processing Workshop_ , 2020. [Online]. Available: https://arxiv.org/pdf/2011.11588.pdf
  * [10] E. Dunbar, M. Bernard, N. Hamilakis, T. A. Nguyen, M. De Seyssel, P. Rozé, M. Rivière, E. Kharitonov, and E. Dupoux, “The Zero Resource Speech Challenge 2021: Spoken language modelling,” in _Interspeech 2021 - Conference of the International Speech Communication Association_ , Brno, Czech Republic, Aug. 2021. [Online]. Available: https://hal.inria.fr/hal-03329301
  * [11] B. van Niekerk, L. Nortje, M. Baas, and H. Kamper, “Analyzing speaker information in self-supervised models to improve zero-resource speech processing,” 08 2021, pp. 1554–1558.
  * [12] A. Baevski, M. Auli, and A. Mohamed, “Effectiveness of self-supervised pre-training for ASR,” in _2020 IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2020, Barcelona, Spain, May 4-8, 2020_. IEEE, 2020, pp. 7694–7698. [Online]. Available: https://doi.org/10.1109/ICASSP40776.2020.9054224
  * [13] A. Baevski, S. Schneider, and M. Auli, “vq-wav2vec: Self-supervised learning of discrete speech representations,” in _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net, 2020. [Online]. Available: https://openreview.net/forum?id=rylwJxrYDS
  * [14] S. Schneider, A. Baevski, R. Collobert, and M. Auli, “wav2vec: Unsupervised pre-training for speech recognition,” in _Interspeech 2019, 20th Annual Conference of the International Speech Communication Association, Graz, Austria, 15-19 September 2019_ , G. Kubin and Z. Kacic, Eds. ISCA, 2019, pp. 3465–3469. [Online]. Available: https://doi.org/10.21437/Interspeech.2019-1873
  * [15] K. Lakhotia, E. Kharitonov, W.-N. Hsu, Y. Adi, A. Polyak, B. Bolte, T.-A. Nguyen, J. Copet, A. Baevski, A. Mohamed, and E. Dupoux, “On Generative Spoken Language Modeling from Raw Audio,” _Transactions of the Association for Computational Linguistics_ , Feb. 2021. [Online]. Available: https://hal.inria.fr/hal-03329219
  * [16] J. Shen, R. Pang, R. J. Weiss, M. Schuster, N. Jaitly, Z. Yang, Z. Chen, Y. Zhang, Y. Wang, R. Skerrv-Ryan, R. A. Saurous, Y. Agiomvrgiannakis, and Y. Wu, “Natural TTS synthesis by conditioning wavenet on MEL spectrogram predictions,” in _2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2018, pp. 4779–4783.
  * [17] V. Panayotov, G. Chen, D. Povey, and S. Khudanpur, “Librispeech: An ASR corpus based on public domain audio books,” in _2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2015, pp. 5206–5210.
  * [18] A. van den Oord, Y. Li, and O. Vinyals, “Representation learning with contrastive predictive coding,” _CoRR_ , vol. abs/1807.03748, 2018. [Online]. Available: http://arxiv.org/abs/1807.03748
  * [19] M. Rivière, A. Joulin, P.-E. Mazaré, and E. Dupoux, “Unsupervised pretraining transfers well across languages,” in _ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing_. Barcelona / Virtual, Spain: IEEE, May 2020, pp. 7414–7418. [Online]. Available: https://hal.archives-ouvertes.fr/hal-02959418
  * [20] M. Ott, S. Edunov, A. Baevski, A. Fan, S. Gross, N. Ng, D. Grangier, and M. Auli, “fairseq: A fast, extensible toolkit for sequence modeling,” in _Proceedings of NAACL-HLT 2019: Demonstrations_ , 2019. [Online]. Available: https://aclanthology.org/N19-4009
  * [21] A. Wang and K. Cho, “BERT has a mouth, and it must speak: BERT as a Markov random field language model,” in _Proceedings of the Workshop on Methods for Optimizing and Evaluating Neural Language Generation_. Minneapolis, Minnesota: Association for Computational Linguistics, Jun. 2019, pp. 30–36. [Online]. Available: https://aclanthology.org/W19-2304
  * [22] J. Salazar, D. Liang, T. Q. Nguyen, and K. Kirchhoff, “Masked language model scoring,” in _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_. Online: Association for Computational Linguistics, Jul. 2020, pp. 2699–2712. [Online]. Available: https://aclanthology.org/2020.acl-main.240
  * [23] E. Kharitonov, J. Copet, K. Lakhotia, T. A. Nguyen, P. Tomasello, A. Lee, A. Elkahky, W.-N. Hsu, A. Mohamed, E. Dupoux, and Y. Adi, “textless-lib: a library for textless spoken language processing,” 2022.

##  Appendix A Appendix

###  A-A Hyper-parameters of trained models

Table VI shows the hyperparameters of all trained models in the paper.

###  A-B Unit-Phone alignments for CPC Discrete Units

We analyse to what extent the discrete units obtained with different numbers
of clusters correlate with the gold phonemes. Using the phoneme alignments of
Librispeech available from [9], we collect all unit-phoneme pairs from the
utterances of the dev-clean subset and compute the probability of each phoneme
given a discrete unit. Figure 3 (top, middle, bottom) shows this unit-phone
alignment for discrete units obtained from CPC features with 20, 50 and 500
clusters respectively. The phoneme order is obtained by clustering the rows of
the 50-unit model with a hierarchical clustering method.

###  A-C Layer-wise analysis of BERT models on ABX metrics

In addition to performing ABX on input and target features of the BERT models,
we also compute the ABX error on the features extracted from hidden
Transformer layers of the trained BERT/HuBERT models. Table VII reports the
average (within-across) ABX errors on the Librispeech dev-clean subset for all
the trained models in the paper.

We see that well-trained models with good language modeling scores (models
1,2,13,15,17,8,9,10; cf. Table III) seem to have very good ABX errors compared
to the others. By looking at the best hidden features of each model, we see
that models with discrete targets (models 1,2,13) are able to reconstruct
hidden features which are better than the targets, while this is not the case
for most models with continuous targets (except model 10).

Figure 3: Probability that each discrete unit belongs to possible phonemes
$P\left(phoneme\mid unit\right)$ for discrete units obtained by clustering CPC
features with different numbers of clusters: 20 (top), 50 (middle) and 500
(bottom). Unit-Phoneme alignments are collected on Librispeech dev-clean
subset. The phoneme order is obtained by clustering the rows of the 50-unit
model with a hierarchical clustering method.

Training hyperparameters Inference hyperparameters n units (if discrete) feat.
masking num prob. est. SIMI librispeech SIMI synthetic id input target stride
length prob. updates $M$ $\Delta t$ layer pooling layer pooling BERT Models on
CPC-big Features 1 50 50 10ms 10 0.5 250k 15 5 11 max 1 min 2 50 50 10ms 10
0.5 250k 15 5 5 min 7 mean 11 50 50 10ms 10 0.5 250k 15 5 11 mean 10 min 12 50
50 10ms 10 0.5 250k 15 5 9 mean 1 min 3 - 50 10ms 10 0.5 250k 35 5 6 min 4 max
4 - 50 10ms 10 0.5 250k 35 5 8 mean 11 min 13 - 50 10ms 10 0.5 250k 25 5 10
max 10 min 14 - 50 10ms 10 0.5 250k 35 5 4 max 6 mean 5 - - 10ms 10 0.5 250k
45 5 8 max 1 mean 15 - - 10ms 10 0.5 250k 35 5 11 max 11 mean 16 - - 10ms 10
0.5 250k 45 5 12 mean 12 mean 6 - - 10ms 10 0.5 250k 35 5 1 mean 12 mean 17 -
- 10ms 10 0.5 250k 25 5 12 min 12 max 18 - - 10ms 10 0.5 250k 35 5 1 mean 12
mean 7 - - 10ms 10 0.5 250k 35 5 3 mean 1 mean 8 50 - 10ms 10 0.5 250k 25 5 12
mean 5 min 19 50 - 10ms 10 0.5 250k 25 5 8 mean 4 min 20 50 - 10ms 10 0.5 250k
25 5 8 mean 11 mean 9 50 - 10ms 10 0.5 250k 15 5 12 max 12 min 21 50 - 10ms 10
0.5 250k 15 5 11 mean 7 min 22 50 - 10ms 10 0.5 250k 15 5 4 mean 12 max 10 50
- 10ms 10 0.5 250k 15 5 10 max 12 max HuBERT Base Models 23 - 100 20ms 10 0.65
250k 15 5 2 max 5 min 24 - 500 20ms 10 0.65 400k 15 5 1 mean 9 max 25 - 500
20ms 10 0.65 400k 15 5 10 mean 6 max BERT Models on HuBERT Discrete Units 26 -
500 20ms 10 0.5 250k 10 5 7 mean 8 mean

TABLE VI:  Hyperparameters of all trained models.

id input feature target feature loss input layer 3 layer 6 layer 9 layer 12
target BERT Models on CPC-big Features 1 CPC-l2+km50 CPC-l2+km50 NLL-l
red!2.95455!white7.3 black!35!green!1.66667!white!5.90 red!0.29545!white6.13
black!35!green!6!white!5.64 black!35!green!35.5!white!3.87
red!2.95455!white7.3 2 CPC-l2+km50 CPC-l2+km50 NLL-e red!2.95455!white7.3
black!35!green!10.66667!white!5.36 red!2.20455!white6.97
black!35!green!10.0!white!5.40 black!35!green!34.16667!white!3.95
red!2.95455!white7.3 11 CPC-l0+km50 CPC-l2+km50 NLL-e red!45.70455!white26.11
red!20.63636!white15.08 red!14.0!white12.16 red!10.52273!white10.63
red!2.29545!white7.01 red!2.95455!white7.3 12 CPC-l4+km50 CPC-l2+km50 NLL-e
red!36.59091!white22.1 red!17.20455!white13.57 red!7.72727!white9.40
red!9.20455!white10.05 red!0.13636!white6.06 red!2.95455!white7.3 3 CPC-l2
CPC-l2+km50 NLL-l black!35!green!36.5!white!3.81 red!6.52273!white8.87
red!21.13636!white15.30 red!20.25!white14.91 red!7.04545!white9.10
red!2.95455!white7.3 4 CPC-l2 CPC-l2+km50 NLL-e black!35!green!36.5!white!3.81
red!5.61364!white8.47 red!24.59091!white16.82 red!35.15909!white21.47
red!15.54545!white12.84 red!2.95455!white7.3 13 CPC-l0 CPC-l2+km50 NLL-e
red!20.47727!white15.01 red!1.29545!white6.57 black!35!green!21!white!4.74
black!35!green!32.66667!white!4.04 black!35!green!32.16667!white!4.07
red!2.95455!white7.3 14 CPC-l4 CPC-l2+km50 NLL-e red!8.52273!white9.75
red!2.34091!white7.03 red!6.31818!white8.78 red!10.20455!white10.49
black!35!green!17!white!4.98 red!2.95455!white7.3 5 CPC-l2 CPC-l2 NCE
black!35!green!36.5!white!3.81 red!1.06818!white6.47 red!1.88636!white6.83
red!0.20455!white6.09 black!35!green!0.33333!white!5.98
black!35!green!36.5!white!3.81 15 CPC-l0 CPC-l2 NCE red!20.47727!white15.01
red!3.29545!white7.45 black!35!green!6.16667!white!5.63
black!35!green!21.16667!white!4.73 black!35!green!1.16667!white!5.93
black!35!green!36.5!white!3.81 16 CPC-l4 CPC-l2 NCE red!8.52273!white9.75
red!1.86364!white6.82 red!2.70455!white7.19 black!35!green!9.0!white!5.46
black!35!green!8.33333!white!5.50 black!35!green!36.5!white!3.81 6 CPC-l2
CPC-l2 L1 black!35!green!36.5!white!3.81 red!3.20455!white7.41
red!4.11364!white7.81 red!9.79545!white10.31 red!15.61364!white12.87
black!35!green!36.5!white!3.81 17 CPC-l0 CPC-l2 L1 red!20.47727!white15.01
red!1.45455!white6.64 black!35!green!21.16667!white!4.73
black!35!green!21.5!white!4.71 black!35!green!8.66667!white!5.48
black!35!green!36.5!white!3.81 18 CPC-l4 CPC-l2 L1 red!8.52273!white9.75
black!35!green!1.83333!white!5.89 black!35!green!12.16667!white!5.27
black!35!green!20!white!4.80 black!35!green!10.33333!white!5.38
black!35!green!36.5!white!3.81 7 CPC-l2 CPC-l2 L2
black!35!green!36.5!white!3.81 red!0.79545!white6.35 red!1.31818!white6.58
red!3.75!white7.65 red!7.27273!white9.20 black!35!green!36.5!white!3.81 8
CPC-l2+km50 CPC-l2 NCE red!2.95455!white7.3 red!1.11364!white6.49
black!35!green!8.33333!white!5.50 red!1.97727!white6.87
black!35!green!10.5!white!5.37 black!35!green!36.5!white!3.81 19 CPC-l0+km50
CPC-l2 NCE red!45.70455!white26.11 red!22.56818!white15.93
red!15.56818!white12.85 red!10.34091!white10.55 red!7.02273!white9.09
black!35!green!36.5!white!3.81 20 CPC-l4+km50 CPC-l2 NCE
red!36.59091!white22.1 red!15.63636!white12.88 red!9.47727!white10.17
red!9.5!white10.18 red!6.38636!white8.81 black!35!green!36.5!white!3.81 9
CPC-l2+km50 CPC-l2 L1 red!2.95455!white7.3 black!35!green!12.5!white!5.25
black!35!green!4.83333!white!5.71 black!35!green!18.5!white!4.89
red!11.27273!white10.96 black!35!green!36.5!white!3.81 21 CPC-l0+km50 CPC-l2
L1 red!45.70455!white26.11 red!16.20455!white13.13 red!11.06818!white10.87
red!8.65909!white9.81 red!21.36364!white15.40 black!35!green!36.5!white!3.81
22 CPC-l4+km50 CPC-l2 L1 red!36.59091!white22.1 red!12.36364!white11.44
red!8.88636!white9.91 red!3.75!white7.65 red!13.36364!white11.88
black!35!green!36.5!white!3.81 10 CPC-l2+km50 CPC-l2 L2 red!2.95455!white7.3
black!35!green!9.66667!white!5.42 black!35!green!15.16667!white!5.09
black!35!green!30!white!4.20 black!35!green!38.66667!white!3.68
black!35!green!36.5!white!3.81 HuBERT Base Models 23 waveform MFCC+km100 NLL-e
- red!2.56818!white7.13 black!35!green!30.33333!white!4.18
black!35!green!16.16667!white!5.03 red!6.93182!white9.05
red!49.72727!white27.88 24 waveform H1-l6+km500 NLL-e - red!1.88636!white6.83
black!35!green!21.5!white!4.71 black!35!green!26.33333!white!4.42
black!35!green!41.16667!white!3.53 red!2.20455!white6.97 25 waveform
H2-l12+km500 NLL-e - red!1.5!white6.66 black!35!green!26.16667!white!4.43
black!35!green!25.33333!white!4.48 black!35!green!37!white!3.78
red!0.59091!white6.26 BERT Models on HuBERT Discrete Units 26 H2-l12+km500
H2-l12+km500 NLL-e red!0.59091!white6.26 black!35!green!11.33333!white!5.32
red!1.15909!white6.51 red!1.68182!white6.74 black!35!green!26.16667!white!4.43
red!0.59091!white6.26

TABLE VII:  Average (within and across) dev-clean ABX error of input features,
target features and features from different hidden layers of Transformer
model. CPC-lx stands for layer x of CPC-big, Hj-lx stands for layer x of
HuBERT j’th iteration.


# Improving Photometric Redshift Estimates with Training Sample Augmentation

Irene Moskowitz Department of Physics and Astronomy, Rutgers, The State
University of New Jersey, Piscataway, NJ 08854, USA  Eric Gawiser Department
of Physics and Astronomy, Rutgers, The State University of New Jersey,
Piscataway, NJ 08854, USA  John Franklin Crenshaw Department of Physics,
University of Washington, Seattle, WA 98195, USA  DIRAC Institute, University
of Washington, Seattle, WA 98195, USA  Brett H. Andrews Department of Physics
and Astronomy and PITT PACC, University of Pittsburgh, Pittsburgh, PA 15260,
USA  Samuel Schmidt Department of Physics and Astronomy, University of
California, Davis, CA, 95616, USA  The LSST Dark Energy Science Collaboration

###### Abstract

Large imaging surveys will rely on photometric redshifts (photo-$z$’s), which
are typically estimated through machine learning methods. Currently planned
spectroscopic surveys will not be deep enough to produce a representative
training sample for LSST, so we seek methods to improve the photo-$z$
estimates from non-representative training samples. Spectroscopic training
samples for photo-$z$’s are biased towards redder, brighter galaxies, which
also tend to be at lower redshift than the typical galaxy observed by LSST,
leading to poor photo-$z$ estimates with outlier fractions nearly 4 times
larger than for a representative training sample. In this paper, we apply the
concept of training sample augmentation, where we augment non-representative
training samples with simulated galaxies possessing otherwise unrepresented
features. When we select simulated galaxies with (g-z) color, i-band magnitude
and redshift outside the range of the original training sample, we are able to
reduce the outlier fraction of the photo-$z$ estimates by nearly $50\%$ and
the normalized median absolute deviation (NMAD) by $56\%$. When compared to a
fully representative training sample, augmentation can recover nearly $70\%$
of the increase in the outlier fraction and $80\%$ of the increase in NMAD.
Training sample augmentation is a simple and effective way to improve training
samples for photo-$z$’s without requiring additional spectroscopic samples.

††journal: ApJ Letters

##  1 Introduction

Understanding the nature of dark energy is a major open question in cosmology.
Stage-IV dark energy experiments, such as the Vera C. Rubin Observatory’s
Legacy Survey of Space and Time (LSST, Ivezić et al. 2019), Euclid (Laureijs
et al., 2011) and Roman (Akeson et al., 2019), are scheduled to come online in
the coming years.

Imaging surveys will need to obtain redshifts to galaxies, but there will be
too many for spectroscopic redshifts to be feasible. LSST alone is expected to
observe billions of galaxies and will therefore rely on photometric redshifts
(photo-$z$’s). Photo-$z$’s can be estimated through machine learning
algorithms, which learn to associate photometric quantities, such as colors
and magnitudes, with a redshift estimate.

Previous Stage-III dark energy surveys have also used machine learning for
estimating photometric redshifts. The Hyper Suprime-Cam Subaru Strategic
Program (HSC-SSP, Aihara et al. 2017) used DNNz and DEMPz (Hsieh & Yee, 2014)
for the Year 3 cosmology results (Rau et al., 2023; Sugiyama et al., 2023;
Miyatake et al., 2023). Both DNNz and DEMPz are conditional density
estimators. The Dark Energy Survey (DES, Abbott et al. 2018) has used a self-
organizing map (SOMPZ, Myles et al. 2021) for estimating photo-$z$’s. The
Kilo-Degree Survey (KiDS, Heymans et al. 2021) has also used self-organizing
maps for photo-$z$ estimation (Hildebrandt et al., 2021).

Machine learning methods require a training sample of galaxies with both
photometry and spectroscopic redshifts. However, existing spectroscopic
samples are biased towards brighter, redder galaxies than LSST will observe in
general, and these also tend to be at lower redshift than the typical LSST
galaxy. This means that training samples for photo-$z$ estimation will not be
representative of LSST data, leading to poor photo-$z$ estimation for galaxies
with photometry not represented in the training sample. The Dark Energy
Spectroscopic Instrument (DESI, Flaugher & Bebek 2014), along with
spectroscopic redshifts from Euclid, Roman and 4MOST (de Jong et al., 2019),
will alleviate this issue to an extent, but the DESI survey will not be as
deep as LSST; additional spectroscopic redshifts from DESI can not solve the
problem alone. We will need methods to improve the redshift estimation that do
not involve obtaining more spectroscopic redshifts.

One method for improving training samples without obtaining more spectroscopic
redshifts is through data augmentation, which is the process of modifying a
training sample in some way to increase the generality of a machine learning
model (Shorten & Khoshgoftaar, 2019). Data augmentation can be done by
transforming existing training sample data in some way, such as through
rotations or deformations in the case of image recognition (Bloice et al.,
2017), or by generating synthetic data for the training sample (Bird et al.,
2021). Broussard & Gawiser (2021) used this synthetic data generation method
for augmentation to estimate photo-$z$’s.

In this paper, we investigate a slightly different method of augmenting the
training sample by adding galaxies from simulated catalogs to our training
sample. By selecting simulated galaxies with photometry and/or redshifts not
otherwise represented in the training sample, this training sample
augmentation can expand the range of feature space capable of producing good
photo-$z$ estimates, provided the simulated catalog used for augmentation has
reasonable colors. If the simulated catalog is too unrealistic, this will only
create confusion in our model.

Section 2 describes our simulated data, including our stand in for real LSST
data and the simulated catalog used for augmenting the training sample.
Section 3 describes our methodology, including how a realistically non-
representative training sample is created, how we estimate photo-$z$’s, and
the process for augmenting the training sample. Section 4 discusses our
results, and section 5 concludes.

##  2 Simulated Data

###  2.1 DC2

The LSST Dark Energy Science Collaboration (DESC) Data Challenge 2 catalog
(DC2;Abolfathi et al. 2021) is a 300 deg2 area of simulated LSST observations.
The base input for DC2 is the CosmoDC2 galaxy catalog (Korytov et al., 2019),
which is derived from the Outer Rim N-body simulation (Heitmann et al., 2019).
Galaxies were assigned to halos using UniverseMachine (Behroozi et al., 2019)
and GalSampler (Hearin et al., 2020). Complete galaxy properties are generated
with Galacticus (Benson, 2012). Galaxy spectra are constructed from stellar
population spectra computed with fsps (Conroy et al., 2009).

To generate the DC2 catalog, stars, supernovae, strong lenses and AGN are
added to the CosmoDC2 catalog. The object catalogs are passed as inputs to the
image simulation software imSim 111https://github.com/LSSTDESC/Imsim to
generate LSST-like images, which are then processed by the LSST Science
Pipelines.

From the DC2 catalog, we select objects with magnitude $i>17$. We require S/N
$>6$ in the i-band, as well as S/N $>3$ in at least one other band. To
minimize contamination from stars, we select only extended objects. This
extendedness cut does not entirely eliminate stars, but the stellar
contamination is low, and the training sample selection process (described in
section 3.1) ends up placing all stars in the application sample.

The base, unaugmented training sample and the application sample of galaxies,
for which we estimate redshifts, are formed from our selected DC2 objects. In
this work, DC2 is a stand-in for real data.

###  2.2 Buzzard

The Buzzard simulation (DeRose et al., 2019) was built from the L-GADGET2 dark
matter simulations (Springel, 2005). Galaxies are assigned to halos using
AddGals (Wechsler et al., 2022) using the abundance matching technique. Galaxy
SEDs were assigned to match the measured SED-luminosity-density relationship
in SDSS data.

We use the Buzzard catalog selected for the DESC Tomographic Challenge (Zuntz
et al., 2021). Details on selection cuts, post-processing and uncertainty
generation can be found in that paper.

The Buzzard method of assigning spectra is completely independent from fsps,
so Buzzard SEDs should be sufficiently different from DC2 SEDs to simulate
adding simulated galaxies to training samples of real galaxies. In this work,
we use the Buzzard catalog as a simulated catalog with which to augment the
DC2 training sample.

##  3 Methodology

###  3.1 Non-representative Training Sample

Figure 1: The results of partitioning our DC2 catalog into training (left) and
application (center) samples. The training sample is redder and brighter than
the bulk of the application sample. The right panel shows the Buzzard sample
used for augmenting the training sample. The horizontal dotted line shows the
i-band selection criterion, while the vertical dotted line shows the (g-z)
color criterion. The dashed line indicates the section criterion for
color+magnitude augmentation, which generally matches the shape of the DC2
training sample in the left panel. Open arrows indicate which regions of
color-magnitude space are used for single feature augmentation, while solid
arrows indicate regions used for color+magnitude augmentation See section 3.3
for more details on augmentation criteria.

Existing spectroscopic galaxy samples are brighter and redder than expected
LSST observations, and also tend to be at lower redshifts. To partition our
DC2 catalog into a realistically non-representative training sample and
application sample, we use the GridSelection degrader in the DESC
RAIL222https://github.com/LSSTDESC/rail software (Schmidt et al., 2023). We
briefly summarize the GridSelection degrader below. A more detailed discussion
can be found in Moskowitz et al. (2023).

The GridSelection degrader is based on the second data release of the Hyper
Suprime Cam Subaru Strategic Program (Aihara et al., 2019). Galaxies with
similar photometry to early LSST observations are selected from the HSC Wide
catalog; some of these galaxies have photometry only, and some have matched
spectroscopic redshifts. The range in i-band magnitude and (g-z) color are
divided into 100x100 pixels. Within each pixel, a ratio of the number of
galaxies with spectroscopic redshifts to the total number of galaxies is
computed, along with the 99th percentile in spectroscopic redshift, denoted
$z_{max}$. The GridSelection degrader divides our DC2 galaxies into the same
set of pixels in i vs (g-z) and automatically assigns DC2 objects with
$z_{true}>z_{max}$ to the application sample. From the remaining DC2 objects,
the GridSelection degrader randomly selects objects for the training sample
such that the ratio of DC2 training objects to total objects in a pixel
matches the ratio from HSC.

After partitioning the full DC2 sample into training and application samples,
the training sample contains 186,837 galaxies, while the application sample
contains 5,520,458 objects. The left and center panels of Figure 1 show the
resulting DC2 training and application samples, where it is clear that the
training sample is redder and brighter than the majority of the application
sample. Figure 2 shows the (normalized) redshift distributions of both
samples, as well as the distributions of the Buzzard sample and the best
performing augmentation choice. The training sample is biased towards lower
redshifts than the application sample as a whole. The right panel of Figure 1
shows the Buzzard sample that will be used for augmentation. See section 3.3
for more details.

Figure 2: Normalized redshift distributions of the DC2 training sample (blue
solid line), DC2 application sample (orange dashed line), and Buzzard sample
(green dotted line). The DC2 training sample is biased to lower redshifts than
the application sample. The vertical dashed line indicates the selection
criterion for redshift augmentation, with the arrow indicating the region of
redshift space used for augmentation. The black dot-dashed line shows the
redshift distribution of the best performing, post augmentation training
sample shown in the top right panel of Figure 4.

###  3.2 Photo-z estimation

To estimate photo-$z$’s, we use FlexZBoost (Izbicki & Lee 2017; Dalmasso et
al. 2020) as implemented in RAIL. FlexZBoost is a non-parametric conditional
density estimator for redshifts. It takes as inputs the magnitudes and errors
in each of the bands ugrizy and outputs a photo-z pdf for each object in the
application sample.

To evaluate the quality of a set of photo-z estimates, we use the outlier
fraction, normalized median absolute deviation (NMAD) and bias. We define an
outlier as $|z_{true}-z_{phot}|/(1+z_{true})>0.15$. The NMAD is given by:

$1.4826\times\textrm{Med}\left[\lvert\frac{\Delta
z}{1+z_{true}}-\textrm{Med}\left(\frac{\Delta
z}{1+z_{true}}\right)\rvert\right],$ (1)

where the bias is given by Median$(\Delta z/(1+z_{true}))$. We use the mean of
the pdf as a point estimate for each photo-$z$.

The photo-$z$’s estimated from the base, non-representative DC2 training
sample are shown in the top panel of Figure 3. The outlier fraction is quite
high at nearly 50%. In particular, the majority of galaxies with
$z_{true}\gtrsim 1.0$ have outlier $z_{phot}$ estimates. This is due to the
fact that our DC2 training sample has very few objects with $z>1.0$. For
comparison, the bottom panel of Figure 3 shows the results for a fully
representative DC2 training sample, which obtains an outlier fraction of 0.14.
This represents the best we can expect to do using FlexZBoost. As shown in
Figure 3, the unaugmented, non-representative training sample produces much
worse photo-$z$’s than the representative training sample, particularly at
$z_{true}>1.0$.

Figure 3: Top: Photo-$z$’s estimated from the realistic, non-representative
DC2 training sample shown in Figure 1. Solid black lines indicate the boundary
for outliers. Bottom: Photo-$z$’s estimated from a fully representative
training sample drawn randomly from the DC2 application sample in Figure 1.

###  3.3 Augmentation

We augment the training sample by adding 10,000 Buzzard galaxies according to
a set of criteria, taking care to only use knowledge about the application
sample that would be available for real data. The simplest criterion for
augmentation is to select Buzzard galaxies with higher redshifts than those
present in the DC2 training sample. We refer to this as redshift augmentation,
and make the selection $z_{buzzard}>1.0$. This region is indicated by the
vertical dashed line and arrow in Figure 2.

Since DC2 application galaxies are also dimmer and bluer than the training
sample, we also choose magnitude and color selection criteria, which we call
magnitude augmentation and color augmentation, respectively. For magnitude
augmentation we make the selection $i_{buzzard}>23$, and for color
augmentation we choose $(g-z)_{buzzard}<1.75$. These boundaries were chosen to
match where the magnitude and (g-z) color distributions in the training sample
start to decline. They are indicated by the dotted lines and open arrows in
the right panel of Figure 1.

We test augmentation with each of the features individually, as well as in
combination with each other.

###  3.4 Photometry Shifts

Since DC2 and Buzzard rely on different methods for determining galaxy SEDs,
the colors as a function of redshift are different between the two
simulations. If this augmentation method was used for real data, it would be
advantageous to shift the simulated photometry to look like the real
photometry in the application sample. Therefore, we also attempt to match the
Buzzard photometry to the DC2 application sample. This modifies the color-
redshift relationship in Buzzard to potentially more closely resemble the
color-redshift relationship of DC2. Since we do not use the true redshifts of
the application sample, this represents something we could do with real data.

####  3.4.1 Magnitude Shifts

The simplest way to transform Buzzard colors is to apply a single shift to the
Buzzard magnitudes to make their median match the median of the DC2
application sample magnitudes in each band. We will refer to this sample as
the ‘magnitude shifted Buzzard’ sample.

In addition to the medians, we can also rescale the NMADs to match in each
band. This is a proxy for matching the first and second moments of the
photometry distributions. To shift the NMADS, we apply the following
transformation:

$mag_{j,new}=\frac{NMAD_{j,DC2}}{NMAD_{j}}\times[mag_{j}-med_{j}]+med_{j}$ (2)

where $NMAD_{j}$ refers to the NMAD in band $j$, $med_{j}$ is the median
magnitude in band $j$ and all quantities are for Buzzard unless indicated by
the DC2 subscript. We will refer to this sample as the ‘NMAD shifted Buzzard’
sample.

####  3.4.2 Normalizing Flows

The simple shift method is able to match the medians and NMADs of the Buzzard
and DC2 color distributions, but not the shapes of the distributions. To
attempt to more fully match the color distributions, we use normalizing flows
to produce a catalog of DC2-like photometry with Buzzard-like redshifts.

We use the PZFlow333https://jfcrenshaw.github.io/pzflow/ package (Crenshaw et
al., 2023) as implemented in RAIL for training the normalizing flows. We train
two flows: one on DC2 photometry, and one on Buzzard photometry. The DC2 flow
learns the probability distribution function of the DC2 photometry,
p(photometry), while the Buzzard flow is a conditional flow that learns the
probability density function of the redshift given the photometry,
p($z|$photometry). The features used for training are i-band magnitudes and
(u-g), (g-r), (r-i), (i-z) and (z-y) colors. We train 100 epochs for the DC2
flow, and 150 epochs for the Buzzard flow.

Once the flows are trained, we sample from the DC2 flow to make a new catalog
of galaxies with DC2-like photometry. We then sample from the Buzzard flow,
using the new DC2-like photometry as conditions, to generate Buzzard-like
redshifts for our DC2-like photometry. Finally, we use the RAIL LSSTErrorModel
degrader to generate LSST-like errors on the magnitudes. This set of DC2-like
photometry and Buzzard-like redshifts constitutes our flowed catalog from
which we draw galaxies for augmentation. We will refer to this sample as the
“flowed Buzzard” sample.

##  4 Results

Figure 4: The best case photo-$z$ estimation for single, double and triple
feature augmentation (bottom row), with the corresponding augmented training
samples (top row). Left: When only a single feature is used for augmentation,
redshift augmentation produces the best photo-$z$ estimates. The training
sample was augmented with the unshifted Buzzard catalog, which in this case
produced better photo-$z$ statistics than the magnitude shifted or flowed
Buzzard catalog. Center: When a double feature combination is used for
selecting galaxies for augmentation, the combination of color+redshift
produces the best results. The best results for this case came from using the
magnitude shifted Buzzard catalog. Right: The best photo-$z$ estimates were
produced when all three features are combined for selected galaxies for
augmentation. The best results for this case came from using the magnitude
shifted Buzzard catalog.

Table 1 summarizes the outlier fractions and NMADS achieved for each
combination of augmentation features and each Buzzard sample (unshifted,
magnitude shifted and flowed), as well as those achieved for the unaugmented
training sample and a fully representative training sample. Every kind of
augmentation we tested improved the outlier fraction and NMAD of the resulting
photo-$z$’s. Augmentations involving redshift selections performed better than
those without. The magnitude shifted Buzzard sample generally produced better
results than the unshifted or flowed Buzzard samples; however, in the case of
selecting galaxies for augmentation using only a single features, the
unshifted Buzzard catalog produced the best results.

Adding the NMAD shift to the magnitude shift did not hurt, but showed no
improvement over the simple magnitude shift, so we only show results for the
magnitude shifted sample. We also tried combining the magnitude shifted
Buzzard sample with the normalizing flow method, but the standard color flow
worked better.

Since FlexZBoost also takes in photometric errors, we tested the effect of
changing the errors. We tested multiplying the errors by factors of 0.1, 2.0
and $2.0\times(1+z)$. Variations in the outlier fractions and NMADs were
smaller than the variations between augmentation cases.

Each method of augmentation produces an outlier fraction and NMAD for the
resulting photo-$z$’s. In addition to the raw metrics, we also report the
ratio of the augmented results to the unaugmented results. Since the best case
scenario is the fully representative training case, not an outlier
fraction/NMAD of 0, we also report a percent recovery towards the outlier
fraction/NMAD achieved in the representative case. The percent recovery is
calculated as $(X_{unrep}-X_{aug})/(X_{unrep}-X_{rep})$, where $X$ is either
the outlier fraction or NMAD, the subscript ‘aug’ refers to the statistic for
the augmentated training sample, the subscript ‘unaug’ refers to the
unaugmented training sample, and the subscript ‘rep’ refers to the
representative training sample.

The following subsections discuss the results for single feature, double
feature and triple feature augmentations. Table 1 summarizes the results.

###  4.1 Augmentation with Individual Features

When augmenting with a single feature, we choose Buzzard galaxies with either
$z_{true}>1.0$, $i$-mag $>23$ or $(g-z)<1.75$. When using a single feature for
augmentation, the base, unshifted Buzzard catalog produced the lowest outlier
fractions and NMADs.

Of the three features, redshift augmentation produces the best results when a
single feature is used for augmentation, while color performs the worst. The
redshift augmented training sample, using the unshifted Buzzard catalog, and
the resulting photo-$z$ estimates for the redshift augmentation are shown in
the left panel of Figure 4. Results for color and magnitude augmentation are,
as well as for the shifted and flowed Buzzard catalogs are listed in Table 1.

###  4.2 Augmentation with Double Feature Combinations

Unaugmented Samples | Outlier | NMAD | Bias |  |  |  |  |  |
---|---|---|---|---|---|---|---|---|---
| Fraction |  |  |  |  |  |  |  |
Representative | 0.141 | 0.057 | 0.0001 |  |  |  |  |  |
Non-Representative | 0.480 | 0.190 | -0.12 |  |  |  |  |  |
Augmented Samples | Unshifted Buzzard | Magnitude Shifted Buzzard | Flowed Buzzard
| Outlier | NMAD | Bias | Outlier | NMAD | Bias | Outlier | NMAD | Bias
| Fraction |  |  | Fraction |  |  | Fraction |  |
Redshift ($z_{buz}>1.0$) | 0.261 | 0.087 | -0.004 | 0.263 | 0.088 | -0.014 | 0.292 | 0.094 | 0.003
Magnitude ($i_{buz}>23$) | 0.318 | 0.097 | -0.030 | 0.407 | 0.138 | -0.069 | 0.327 | 0.107 | -0.039
Color ($(g-z)_{buz}<1.75$) | 0.324 | 0.099 | -0.031 | 0.401 | 0.134 | -0.066 | 0.319 | 0.102 | -0.035
Magnitude+Redshift | 0.268 | 0.090 | -0.001 | 0.259 | 0.086 | -0.015 | 0.293 | 0.096 | 0.001
Color+Redshift | 0.271 | 0.090 | -0.005 | 0.258 | 0.086 | -0.016 | 0.286 | 0.093 | 0.0004
Color+Magnitude | 0.311 | 0.093 | -0.028 | 0.400 | 0.134 | -0.067 | 0.327 | 0.107 | -0.040
Color+Magnitude+Redshift | 0.268 | 0.089 | -0.002 | 0.245 | 0.084 | -0.014 | 0.284 | 0.092 | 0.002
Table 1: Summary of outlier fractions, NMADs and bias achieved for all color-
shifting and augmentation cases. The bolded values correspond to the best
performing augmentation case.

There are three double feature combinations possible: magnitude+redshift,
color+redshift and color+magnitude. Since the training sample creates a
compact shape in color-magnitude space, we fit two lines to mimic the shape of
the top of the color-magnitude distribution of the training sample. We then
choose objects above and to the left of this region, see the dashed line and
solid arrows in the right panel of Figure 1.

The training sample does not have a simple shape in either color-redshift or
magnitide-redshift space. For these feature combinations, we use the
intersection of the selection requirements for each feature rather than the
union; for example, the magnitude+redshift augmentation selects Buzzard
galaxies with both $i$-mag $>23$ and $z_{true}>1.0$. The intersection worked
mildly better than the union.

When using multiple features to select Buzzard galaxies for augmentation, the
magnitude shifted Buzzard sample produced the best results. Results for the
unshifted and flowed Buzzard catalogs are also listed in Table 1.

Similar to the single feature augmentation, the double feature combinations
that include redshift as a selection criterion perform better than
combinations without redshift. Magnitude+Redshift and Color+Redshift
combinations perform virtually identically for the magnitude shifted Buzzard
catalog, with outlier fractions of 0.259 and 0.258, respectivly. This
corresponds to a ratio to the unaugmented results of 0.54, and a 65% recovery
of the fully representative case. Both produce an NMAD of 0.086, which
corresponds to a ratio to the unaugmented results of 0.452 and a 78% recovery
towards the fully representative case. The color+redshift augmented training
sample, using the magnitude shifted Buzzard catalog, and the resulting
photo-$z$ estimates are shown in the middle panel of Figure 4.

When compared to the outlier fraction for the redshift augmentation using the
magnitude shifted catalog (see Table 1), adding color or magnitude provides a
small improvement. The color+magnitude combination also provides a small
improvement over either color or magnitude alone in most cases.

###  4.3 Color+Magnitude+Redshift Augmentation

Finally, we use all three features to select Buzzard galaxies for
augmentation. We use the intersection of the Color+Magnitude and redshift
selection criteria as this provides better results than the union. This is
likely because we always augment the training sample with 10,000 galaxies; in
this case, the intersection most efficiently probes the feature space not
covered by the DC2 training sample. The color+magnitude+redshift augmented
training sample, using the shifted magnitude Buzzard catalog, is shown in the
top right of Figure 4, and the post augmentation redshift distribution is
shown as the black dot-dashed line in Figure 2.

As with the double feature augmentation cases, the simple magnitude shifted
Buzzard catalog produced the best results, and we show that case in the right
panel of Figure 4.

This augmentation case produces the lowest outlier fraction and NMAD out of
all combinations of color shifting and augmentation features. At an outlier
fraction of 0.245, which is a ratio of 0.51 to the unaugmented case, and
recovers $69\%$ of the increase resulting from the non-representative training
sample. We achieve an NMAD of 0.084, a ratio of 0.44 to the unaugmented case,
and an $80\%$ recovery of the increase in NMAD compared to the fully
representative training sample. We also achieve 8 times less bias with this
augmentation than in the unaugmented case (see Table 1).

It can be seen that the augmented training samples using the shifted
magnitudes Buzzard catalog have i-band magnitudes that extend fainter than the
application sample. We tried imposing an $i<26$ cut on the magnitude shifted
Buzzard catalog before selecting galaxies for augmenting, but this produces
slightly worse photo-$z$ estimates. We suspect this is because FlexZBoost is a
conditional density estimator, and a hard cut-off in the magnitudes makes it
difficult for FlexZBoost to estimate the density at the cut-off magnitude.
Allowing the magnitudes to drop off naturally makes it easier for FlexZBoost
to estimate the density at $i=26$, resulting in better photo-$z$ estimates
even though the magnitude range extends farther than required to match the
application sample.

##  5 Conclusions

Large imaging surveys, such as LSST, will not have access to representative
training samples for estimating photometric redshifts. When estimating
photo-$z$’s using a realistically non-representative training sample, the
outlier fraction reaches nearly $50\%$, almost a factor of 3.5 worse than when
using a representative training sample, and a similar increase in scatter as
probed by the NMAD. Obtaining new spectroscopic samples of dim galaxies cannot
solve this problem alone, as it is not feasible to obtain a large enough
sample by the time LSST is expected to see first light. Training sample
augmentation is an easy way to improve photo-$z$ estimates without requiring
additional spectroscopic samples of dim galaxies.

We used the DESC DC2 simulation as a stand-in for eventual LSST data, and
investigated how augmenting a realistically non-representative training sample
with simulated galaxies from Buzzard can improve the photo-$z$ estimates. Even
a relatively simple augmentation process of selecting simulated galaxies with
redshifts higher than those present in the training sample can recover $65\%$
of the increase in the outlier fraction when compared to a fully
representative training sample.

Shifting the photometry of the galaxy catalog used for augmenting the training
sample can improve the results further. We shifted all Buzzard magnitudes so
the median magnitude in each band matches the median magnitudes of the DC2
application sample. We then select galaxies for augmentation in regions of
color-magnitude-redshift space not covered by the training sample. The
resulting photo-$z$ estimates, shown in the right panel of Figure 4 and Table
1, have an outlier fraction below $25\%$, NMAD of 0.084, and bias of -0.014,
representing a nearly $50\%$ reduction over the unaugmented photo-$z$ outlier
fraction, $56\%$ reduction in NMAD, and factor of 8 reduction in bias.
Augmentation has recovered $69\%$ of the increase in outlier fraction compared
to the fully representative case, $80\%$ of the increase in NMAD, and $88\%$
of the increase in bias. With these results in mind, it is clear that training
sample augmentation should be considered in the photo-$z$ pipeline for large
galaxy surveys, including LSST.

Since redshift seems to be the most important feature for augmentation, the
Buzzard catalog is not the best-case scenario for augmentation given how few
galaxies there are at $z>1.0$. There are almost no Buzzard objects at $z>1.5$,
while there are still many application sample objects in that redshift range.
When using augmentation for real data, choosing a simulation with sufficient
redshift coverage, such as the DC2 simulation should produce better results
than were achieved here.

While we leave extensions of this method to real data to future work, we
discuss below a few possible avenues to explore training sample augmentation
for real data. One could, for example, use the deep photo-z catalogs from
COSMOS2020 (Weaver et al., 2022) as a truth catalog, where the COSMOS2020
photo-z catalog would take the place of the DC2 catalog used in this work. It
could also be worthwile to explore using a deep photo-z catalog such as
COSMO2020 as the augmentation catalog, in place of the Buzzard catalog.

Self-organizing maps (SOMs) have also been used to correct for spectroscopic
incompleteness of training samples in DES (see, for example, Myles et al.
2021; Buchs et al. 2019) and KiDS (van den Busch et al., 2022). Assigning
training and application samples to SOM cells can be useful for identifying
under-represented regions of photometry in the training sample that can be
targeted for spectroscopic followup (Masters et al., 2015). Future efforts
could use the SOM methodology to identify photometry regions to target with
augmented training samples.

We have assumed in this work that the outlier fraction and NMAD are good
indicators of photo-$z$ quality. While this is a reasonable assumption, a full
quantification of the improvements provided by augmentation would come from
using the photo-$z$ estimates to do a full cosmological parameter estimation.
For augmentation to truly be useful, it should result in better cosmological
parameter estimates than the unaugmented photo-$z$’s. This analysis will be
presented in a forthcoming paper.

##  6 Acknowledgements

This paper has undergone internal review by the LSST Dark Energy Science
Collaboration. The internal reviewers were Huan Lin and Eve Kovacs. The
authors thank the internal reviewers for their valuable comments.

IM and EG acknowledge support for this research from the LSST Corporation via
grant #2021-42. IM and EG also acknowledge support from the U.S. Department of
Energy, Office of Science, Office of High Energy Physics Cosmic Frontier
Research program under Award Number DE-SC0010008. JFC acknowledges support
from the U.S. Department of Energy, Office of Science, Office of High Energy
Physics Cosmic Frontier Research program under Award Number DE-SC0011665. BHA
acknowledges support by the National Science Foundation under Award Number
AST-2009251.

Author contributions are as follows. IM performed the analysis and wrote the
majority of the paper. EG advised IM, suggested approaches, and provided
feedback on the text. JFC suggested and advised on methodology for the
normalizing flow, and provided feedback on manuscript. BHA provided guidance
on bias discussions and feedback on the manuscript. SS advised on the use of
FlexZBoost and feedback on the manuscript.

The DESC acknowledges ongoing support from the Institut National de Physique
Nucléaire et de Physique des Particules in France; the Science & Technology
Facilities Council in the United Kingdom; and the Department of Energy, the
National Science Foundation, and the LSST Corporation in the United States.
DESC uses resources of the IN2P3 Computing Center (CC-IN2P3–Lyon/Villeurbanne
- France) funded by the Centre National de la Recherche Scientifique; the
National Energy Research Scientific Computing Center, a DOE Office of Science
User Facility supported by the Office of Science of the U.S. Department of
Energy under Contract No. DE-AC02-05CH11231; STFC DiRAC HPC Facilities, funded
by UK BEIS National E-infrastructure capital grants; and the UK particle
physics grid, supported by the GridPP Collaboration. This work was performed
in part under DOE Contract DE-AC02-76SF00515.

## References

  * Abbott et al. (2018) Abbott, T., Abdalla, F., Alarcon, A., et al. 2018, Physical Review D, 98, doi: 10.1103/physrevd.98.043526
  * Abolfathi et al. (2021) Abolfathi, B., Alonso, D., Armstrong, R., et al. 2021, The Astrophysical Journal Supplement Series, 253, 31, doi: 10.3847/1538-4365/abd62c
  * Aihara et al. (2017) Aihara, H., Arimoto, N., Armstrong, R., et al. 2017, Publications of the Astronomical Society of Japan, 70, doi: 10.1093/pasj/psx066
  * Aihara et al. (2019) Aihara, H., AlSayyad, Y., Ando, M., et al. 2019, PASJ, 71, 114, doi: 10.1093/pasj/psz103
  * Akeson et al. (2019) Akeson, R., Armus, L., Bachelet, E., et al. 2019, The Wide Field Infrared Survey Telescope: 100 Hubbles for the 2020s, arXiv, doi: 10.48550/ARXIV.1902.05569
  * Behroozi et al. (2019) Behroozi, P., Wechsler, R. H., Hearin, A. P., & Conroy, C. 2019, Monthly Notices of the Royal Astronomical Society, 488, 3143–3194, doi: 10.1093/mnras/stz1182
  * Benson (2012) Benson, A. J. 2012, New Astronomy, 17, 175–197, doi: 10.1016/j.newast.2011.07.004
  * Bird et al. (2021) Bird, J., Pritchard, M., Fratini, A., Ekárt, A., & Faria, D. 2021, IEEE Robotics and Automation Letters, PP, 1, doi: 10.1109/LRA.2021.3056355
  * Bloice et al. (2017) Bloice, M. D., Stocker, C., & Holzinger, A. 2017, CoRR, abs/1708.04680
  * Broussard & Gawiser (2021) Broussard, A., & Gawiser, E. 2021, The Astrophysical Journal, 922, 153, doi: 10.3847/1538-4357/ac2147
  * Buchs et al. (2019) Buchs, R., Davis, C., Gruen, D., et al. 2019, MNRAS, 489, 820, doi: 10.1093/mnras/stz2162
  * Conroy et al. (2009) Conroy, C., Gunn, J. E., & White, M. 2009, ApJ, 699, 486, doi: 10.1088/0004-637X/699/1/486
  * Crenshaw et al. (2023) Crenshaw, J. F., Yan, Z., & Doster, V. 2023, jfcrenshaw/pzflow: v3.1.1, v3.1.1, Zenodo
  * Dalmasso et al. (2020) Dalmasso, N., Pospisil, T., Lee, A. B., et al. 2020, Astronomy and Computing, 30, 100362, doi: 10.1016/j.ascom.2019.100362
  * de Jong et al. (2019) de Jong, R. S., Agertz, O., Berbel, A. A., et al. 2019, The Messenger, 175, 3, doi: 10.18727/0722-6691/5117
  * DeRose et al. (2019) DeRose, J., Wechsler, R. H., Becker, M. R., et al. 2019, The Buzzard Flock: Dark Energy Survey Synthetic Sky Catalogs.  https://arxiv.org/abs/1901.02401
  * Flaugher & Bebek (2014) Flaugher, B., & Bebek, C. 2014, in Society of Photo-Optical Instrumentation Engineers (SPIE) Conference Series, Vol. 9147, Ground-based and Airborne Instrumentation for Astronomy V, ed. S. K. Ramsay, I. S. McLean, & H. Takami, 91470S, doi: 10.1117/12.2057105
  * Hearin et al. (2020) Hearin, A., Korytov, D., Kovacs, E., et al. 2020, Monthly Notices of the Royal Astronomical Society, 495, 5040–5051, doi: 10.1093/mnras/staa1495
  * Heitmann et al. (2019) Heitmann, K., Finkel, H., Pope, A., et al. 2019, ApJS, 245, 16, doi: 10.3847/1538-4365/ab4da1
  * Heymans et al. (2021) Heymans, C., Tröster, T., Asgari, M., et al. 2021, Astronomy & Astrophysics, 646, A140, doi: 10.1051/0004-6361/202039063
  * Hildebrandt et al. (2021) Hildebrandt, H., van den Busch, J. L., Wright, A. H., et al. 2021, A&A, 647, A124, doi: 10.1051/0004-6361/202039018
  * Hsieh & Yee (2014) Hsieh, B. C., & Yee, H. K. C. 2014, ApJ, 792, 102, doi: 10.1088/0004-637X/792/2/102
  * Ivezić et al. (2019) Ivezić, Ž., Kahn, S. M., Tyson, J. A., et al. 2019, ApJ, 873, 111, doi: 10.3847/1538-4357/ab042c
  * Izbicki & Lee (2017) Izbicki, R., & Lee, A. B. 2017, Electronic Journal of Statistics, 11, 2800 , doi: 10.1214/17-EJS1302
  * Korytov et al. (2019) Korytov, D., Hearin, A., Kovacs, E., et al. 2019, ApJS, 245, 26, doi: 10.3847/1538-4365/ab510c
  * Laureijs et al. (2011) Laureijs, R., Amiaux, J., Arduini, S., et al. 2011, Euclid Definition Study Report, arXiv, doi: 10.48550/ARXIV.1110.3193
  * Masters et al. (2015) Masters, D., Capak, P., Stern, D., et al. 2015, ApJ, 813, 53, doi: 10.1088/0004-637X/813/1/53
  * Miyatake et al. (2023) Miyatake, H., Sugiyama, S., Takada, M., et al. 2023, Hyper Suprime-Cam Year 3 Results: Cosmology from Galaxy Clustering and Weak Lensing with HSC and SDSS using the Emulator Based Halo Model.  https://arxiv.org/abs/2304.00704
  * Moskowitz et al. (2023) Moskowitz, I., Gawiser, E., Bault, A., et al. 2023, The Astrophysical Journal, 950, 49, doi: 10.3847/1538-4357/accc88
  * Myles et al. (2021) Myles, J., Alarcon, A., Amon, A., et al. 2021, Monthly Notices of the Royal Astronomical Society, 505, 4249, doi: 10.1093/mnras/stab1515
  * Rau et al. (2023) Rau, M. M., Dalal, R., Zhang, T., et al. 2023, Monthly Notices of the Royal Astronomical Society, 524, 5109, doi: 10.1093/mnras/stad1962
  * Schmidt et al. (2023) Schmidt, S., Gschwend, J., Crenshaw, J. F., et al. 2023, LSSTDESC/RAIL: v0.98.5, v0.98.5, Zenodo, doi: 10.5281/zenodo.7927358
  * Shorten & Khoshgoftaar (2019) Shorten, C., & Khoshgoftaar, T. 2019, Journal of Big Data, 6, 60, doi: 10.1186/s40537-019-0197-0
  * Springel (2005) Springel, V. 2005, Monthly Notices of the Royal Astronomical Society, 364, 1105–1134, doi: 10.1111/j.1365-2966.2005.09655.x
  * Sugiyama et al. (2023) Sugiyama, S., Miyatake, H., More, S., et al. 2023, Hyper Suprime-Cam Year 3 Results: Cosmology from Galaxy Clustering and Weak Lensing with HSC and SDSS using the Minimal Bias Model.  https://arxiv.org/abs/2304.00705
  * van den Busch et al. (2022) van den Busch, J. L., Wright, A. H., Hildebrandt, H., et al. 2022, Astronomy & Astrophysics, 664, A170, doi: 10.1051/0004-6361/202142083
  * Weaver et al. (2022) Weaver, J. R., Kauffmann, O. B., Ilbert, O., et al. 2022, The Astrophysical Journal Supplement Series, 258, 11, doi: 10.3847/1538-4365/ac3078
  * Wechsler et al. (2022) Wechsler, R. H., DeRose, J., Busha, M. T., et al. 2022, The Astrophysical Journal, 931, 145, doi: 10.3847/1538-4357/ac5b0a
  * Zuntz et al. (2021) Zuntz, J., Lanusse, F., Malz, A. I., et al. 2021, The Open Journal of Astrophysics, 4, doi: 10.21105/astro.2108.13418