# Balancing utility and cognitive cost in social representation

Max Taylor-Davies
School of Informatics
University of Edinburgh
<EMAIL_ADDRESS>
&Christopher G. Lucas
School of Informatics
University of Edinburgh

###### Abstract

To successfully navigate its environment, an agent must construct and maintain
representations of the other agents that it encounters. Such representations
are useful for many tasks, but they are not without cost. As a result, agents
must make decisions regarding how much information they choose to represent
about the agents in their environment. Using selective imitation as an example
task, we motivate the problem of finding agent representations that optimally
trade off between downstream utility and information cost, and illustrate two
example approaches to resource-constrained social representation.

## Representing agents under cost constraints

It is generally accepted that, in order to produce adaptive behaviour, an
agent must in some sense acquire and maintain an internal representation of
its environment (Craik, 1943; Tolman, 1948; Wilson et al., 2014). Furthermore,
unless it is condemned to an entirely solitary existence, we can expect that
many environments encountered by a hypothetical agent will contain other
agents. Much as our agent should represent the rest of the environment, we
expect that it ought also to be able to represent these other agents. Humans
do this—in fact, it seems we are intrinsically motivated to form mental
representations of the other people we encounter (Dennett, 1987; Malle, 2008;
Baker et al., 2017). We use these representations for a variety of different
purposes: understanding the strengths and weaknesses of a colleague in order
to effectively collaborate with them; determining whether an unfamiliar person
should be treated as friend or foe; or predicting the plays of a chess
opponent in order to beat them. Ideally, we would want to maintain
representations that encode all possible available information about every
agent in our environment. But real agents, whether biological or artificial,
will inevitably have to contend with limits on their cognitive resources.
Inspired by a growing body of work in computational cognitive science that
models human cognition through the lens of resource rationality (Lieder and
Griffiths, 2019; Bhui et al., 2021), we consider the problem of developing
social representations that balance downstream utility against cognitive cost.

## Social decision tasks

In order to consider the idea of using social representations for decision-
making, we develop the general construct of a social decision task—which
describes any decision task where the optimal strategy depends on having
access to representations of some set of other agents. That is, the optimal
policy or choice rule (in terms of maximising expected return) is conditional
on the value of some social representation $\chi$:
$\pi(\cdot)=\pi(\cdot|\chi)$. To consider the utility of a particular social
representation $\chi$, first let $\Pi(\chi)$ be the class of possible policies
that depend on $\chi$ (i.e. may produce a different output depending on the
value of $\chi$). We can then consider the utility of $\chi$ as the expected
return induced by the best member of $\Pi(\chi)$:

$U(\chi)=\max_{\pi\in\Pi(\chi)}\mathbb{E}[R\ |\
\text{behaviour}\sim\pi(\cdot|\chi)]$ (1)

This tells us, given access to social representation $\chi$, how well do we
expect to perform on a particular task, assuming that we’re able to make the
best possible use of the information in $\chi$? Given some function
$C:\chi\to\mathbb{R}$ that measures cognitive cost, we define the cost-
adjusted utility of $\chi$ as

$U^{\prime}(\chi)=(1-\lambda)U(\chi)-\lambda C(\chi)$ (2)

where $\lambda\in[0,1]$ is a parameter that trades off between utility and
cost (essentially telling us how much we _care_ about cost in the context of
the current task/environment). A representation $\chi^{*}$ is then considered
optimal if it satisfies $\chi^{*}=\arg\max_{\chi}U^{\prime}(\chi)$. This
optimality criterion is similar to the objectives used within work on
capacity-limited Bayesian decision-making and RL, such as Arumugam et al.
(2023). The key difference (beyond our explicit focus on social
representations) is that we are interested not so much in the cognitive cost
of converting representations into behaviour, but in the cost the
representations themselves. In general, we expect this be a combination of the
cost involved in _acquiring_ a representation, and the cost involved in
_storing_ it—for now we simply blanket both under entropy (for more details
see Appendix A), with Eq. 2 then becoming

$U^{\prime}(\chi)=(1-\lambda)U(\chi)-\lambda H[\chi]$ (3)

This makes intuitive sense as an objective—essentially saying that we want to
maximise expected utility while minimising the amount of information we store
about other agents.

## Selective imitation as an example downstream task

As a simple example of a social decision task, we consider selective
imitation, a common component of human and animal social behaviour which
involves the problem of identifying the ‘best’ (most rewarding) agent to
imitate within a given environment (Rendell et al., 2011; Heyes, 2016). In the
most general framing of this task, the imitating agent $\alpha^{\text{ego}}$
is situated within some environment populated by a number of other agents,
with the ability to observe their behaviour. At each trial,
$\alpha^{\text{ego}}$ must select an agent $\alpha^{\text{target}}$ to
imitate. All agents then execute some sequence of behaviour within the
environment, visible to $\alpha^{\text{ego}}$. $\alpha^{\text{ego}}$ executes
the exact same behaviour as the chosen agent $\alpha^{\text{target}}$, and
receives some reward corresponding to their own internal value function
$\mathbf{v}^{\text{ego}}$. A general policy for selecting imitation targets on
the basis of representations is given by

$\text{Pr}(\alpha^{\text{target}}=\alpha^{(m)})\propto\exp\bigg{(}\frac{w(\alpha^{(m)};\chi)}{\beta^{\text{ego}}}\bigg{)}$
(4)

where $w$ is some function that assigns agent weights based on the
representations $\chi$.

Figure 1: We generate a number of agents, where each has a randomly sampled
scalar value function $\mathbf{v}^{(m)}\in[0,1]$, which controls the agent’s
relative preference for the two states labelled A and B. From each agent we
sample a number of trajectories using different values of the decision noise
parameter $\beta^{(m)}$. We then plot average imitation reward as a function
of $\mathbf{v}^{(m)}$ and $\beta^{(m)}$.

Of course, the optimal choice of $w$ will depend on the type of decision-
making employed by the other agents. We will assume that each agent attempts
to maximise its own return, governed by its own value function
$\mathbf{v}^{(m)}$, and that it does so in a noisily rational manner,
according to a Boltzmann policy with decision noise $\beta^{(m)}$. Under this
model, the imitation reward should be maximised by selecting an agent such
that we maximise some measure of the _similarity_ between
$\mathbf{v}^{\text{target}}$ and $\mathbf{v}^{\text{ego}}$ while minimising
$\beta^{(m)}$. We can demonstrate this through simulation using a simplified
one-dimensional gridworld environment—Fig. 1 shows that imitation return
indeed increases with increasing value function similarity and decreasing
$\beta^{(m)}$. For the remainder of this paper, we will make the assumption
that all agents are equivalently rational, and thus set aside the decision
noise parameter. If we were to consider only downstream utility and disregard
cost, the optimal social representation for selective imitation would
therefore consist simply of every agent’s exact value function. More
generally, we can consider representations of the form
$\chi^{(m)}=\hat{\mathbf{v}}^{(m)}$ where $\hat{\mathbf{v}}^{(m)}$ is some
approximation of $\mathbf{v}^{(m)}$.

To study strategies for selective imitation, we use a contextual bandit
setting. At each trial, each agent $m$ is placed at a random tile within a 2D
grid, chooses an destination tile to travel to, and then receives a reward
$r=\mathbf{v}^{(m)}_{\text{dest}}-c\big{(}|x_{\text{dest}}-x_{\text{start}}|+|y_{\text{dest}}-y_{\text{start}}|\big{)}$
where $c\geq 0$ is a constant that determines the cost of taking one step on
the grid. Agents are noisily rational in selecting destination tiles to
maximise reward, following a Boltzmann policy with constant decision noise
$\beta$. Following the design of Wu et al. (2018) and Witt et al. (2023), we
generate spatially correlated value functions (i.e. nearby tiles generally
yield similar rewards) by sampling from a Gaussian process prior with a radial
basis function kernel. The imitator selects targets according to Eq. 4, with
weights given by

$w(\alpha^{(m)};\chi)=\text{sim}(\mathbf{v}^{\text{ego}},\hat{\mathbf{v}}^{(m)})=\frac{\mathbf{v}^{\text{ego}}\cdot\hat{\mathbf{v}}^{(m)}}{|\hat{\mathbf{v}}^{(m)}||\hat{\mathbf{v}}^{(m)}|}$
(5)

## Experiment 1: state aggregation

Figure 2: Left: illustration of state-aggregated estimates of an example value
function (lighter grid squares indicate higher values). Middle: average return
from selective imitation using state-aggregated value function estimates as
agent representation (with indiscriminate imitation return as a baseline).
Right: cost-adjusted return for aggregated representations at different values
of $\lambda$. Since we are interested only in the tradeoff between return and
cost (entropy), rather than each quantity’s absolute values, both are
individually normalised to lie in [0,1].

If our imitator $\alpha^{\text{ego}}$ has to choose between only a small
number of agents, or faces only a small state space, then it may be feasible
just to represent value functions exactly (i.e. use
$\hat{\mathbf{v}}^{(m)}=\mathbf{v}^{(m)}$, even for $\lambda>0$. But if the
agent population or state space is large, then the optimal representation in
terms of cost-adjusted utility (Eq. 3) will likely be an approximation that
discards some information for the sake of lower entropy (see Appendix A). One
way we might obtain such an approximation is simply through state
aggregation—i.e. partition the state space into non-overlapping square
‘patches’, and then group all states within each patch under a single value
(Sutton and Barto, 1998). Fig. 2 shows how the cost-balanced utility of
aggregated value function representations under various values of the tradeoff
parameter $\lambda$. As expected, we observe that increasing the amount of
aggregation leads to a decrease in both average return and representation
cost, with the relative advantage or disadvantage determined by the tradeoff
parameter: as $\lambda$ increases, the optimal aggregation amount shifts to
the right.

## Experiment 2: representing social groups

Figure 3: Top: cost-adjusted return for group-only and individual
representations (with indiscriminate imitation baseline) for different values
of the group variance ratio $\rho$ and the tradeoff parameter $\lambda$. As
before, both return and cost (entropy) are normalised to lie in [0,1]. Bottom:
illustration of example value functions sampled for a 3-member group for three
different values of the variance ratio $\rho$.

So we have seen how naive compression of value function information can help
an imitator navigate this tradeoff. But rather than merely taking individual
representations and making them coarser, a more humanlike approach might
involve using information about the agent population _as a whole_ , and how
individual agents within the population relate to one another. For instance,
we often represent people that we’re unfamiliar with in terms of their
membership of certain social _groups_ or _categories_ , using our knowledge of
those groups to make inferences about the unknown properties of individuals
within them (Rhodes and Baron, 2019; Liberman et al., 2017). To see how this
can produce cheap agent representations, imagine that we have a known fixed
number of groups $K$, and a population of $M$ agents that are partitioned
among them such that each agent only belongs to one group (given by
$\mathbf{z}_{m}$). Continuing our example of value-function-guided selective
imitation, we will say that each group $k$ has a corresponding distribution
over value functions, $p(\mathbf{v}^{(m)}|\mathbf{z}_{m}=k)$, with mean value
$\mathbf{\mu}^{(k)}$. As an extreme, if we treat all members of group $k$ has
having $\mathbf{v}=\mathbf{\mu}^{(k)}$, then we need only represent each agent
with a single scalar value $\mathbf{z}_{m}$. Including the group means, this
would yield an overall representation cost

$C(\chi_{\text{grp}})=KH[\mathbf{\mu}^{(k)}]+MH[\mathbf{z}_{m}]$ (6)

Intuitively, the extent to which this representation is useful will be
primarily determined by the ratio between inter- and intragroup variation.
When groups are very different, and agents within each group are very similar,
then knowing an agent’s group assignment tells you a lot; when groups are
similar or agents vary a lot within each group, then group assignments are
less informative. We can test this using the same selective imitation bandit
setup as in the previous experiment. We first sample a number of group mean
value functions $\mathbf{\mu}^{(k)}$ from the original GP prior; for each
group we then sample a number of agents from a multivariate Gaussian
distribution with mean $\mathbf{\mu}^{(k)}$ and covariance
$\Sigma=\rho\Sigma_{\text{GP}}$ where $\rho>0$ is a scalar and
$\Sigma_{\text{GP}}$ is the covariance of the GP prior. Fig. 3 compares the
cost-adjusted utility of the groups-only and individuals-only representation
strategies for different values of $\rho$ and $\lambda$. We can see that,
below a certain $\lambda$ threshold, individual representations are always as
good or better than group representations; above a certain $\lambda$
threshold, the reverse is true. Otherwise, as predicted, the optimal
representation depends on the group variance ratio, with group-only
representations becoming less useful as $\rho$ increases.

## Discussion

In this paper, we have briefly laid out and motivated the problem of cost-
constrained social representation; that is, representing other agents in such
a way that optimally trades off between downstream utility and information
cost. Using selective imitation as an example task, we explored two potential
approaches to this problem: compressing individual representations, and
representing agents in terms of group identities. While we hope the work we
present here has value and interest, it is of course very much a preliminary
step—in future work we hope to consider how agents can _learn_ good social
representations by optimising directly for this tradeoff, as well as
broadening our analysis to encompass additional example tasks and agent
features. As hinted at earlier, we also hope to explore a more detailed
conception of representation cost that considers both acquisition and storage
of information.

## References

  * Arumugam et al. [2023] D. Arumugam, M. K. Ho, N. D. Goodman, and B. V. Roy.  Bayesian reinforcement learning with limited cognitive load.  _ArXiv_ , abs/2305.03263, 2023.  URL https://api.semanticscholar.org/CorpusID:258546889.
  * Baker et al. [2017] C. L. Baker, J. Jara-Ettinger, R. Saxe, and J. B. Tenenbaum.  Rational quantitative attribution of beliefs, desires and percepts in human mentalizing.  _Nature Human Behaviour_ , 1(4):1–10, 2017.
  * Bhui et al. [2021] R. Bhui, L. Lai, and S. J. Gershman.  Resource-rational decision making.  _Current Opinion in Behavioral Sciences_ , 41:15–21, 2021.
  * Cover and Thomas [2006] T. M. Cover and J. A. Thomas.  _Elements of information theory / Thomas M. Cover, Joy A. Thomas._ Wiley-Interscience, Hoboken, N.J, second edition. edition, 2006.  ISBN 0471241954.
  * Craik [1943] K. Craik.  Hypothesis on the nature of thought, 1943.
  * Dennett [1987] D. C. Dennett.  The intentional stance.  1987\.  URL https://api.semanticscholar.org/CorpusID:143977161.
  * Heyes [2016] C. Heyes.  Who knows? metacognitive social learning strategies.  _Trends in Cognitive Sciences_ , 20:204–213, 2016.
  * Liberman et al. [2017] Z. Liberman, A. L. Woodward, and K. D. Kinzler.  The origins of social categorization.  _Trends in Cognitive Sciences_ , 21:556–568, 2017.
  * Lieder and Griffiths [2019] F. Lieder and T. L. Griffiths.  Resource-rational analysis: Understanding human cognition as the optimal use of limited computational resources.  _Behavioral and Brain Sciences_ , 43, 2019.
  * Malle [2008] B. F. Malle.  Chapter 12 - the fundamental tools, and possibly universals, of human social cognition.  In R. M. Sorrentino and S. Yamaguchi, editors, _Handbook of Motivation and Cognition Across Cultures_ , pages 267–296. Academic Press, San Diego, 2008.  ISBN 978-0-12-373694-9.  doi:https://doi.org/10.1016/B978-0-12-373694-9.00012-X.  URL https://www.sciencedirect.com/science/article/pii/B978012373694900012X.
  * Rasmussen and Williams [2006] C. E. Rasmussen and C. K. I. Williams.  _Gaussian processes for machine learning / Carl Edward Rasmussen, Christopher K.I. Williams._ Adaptive computation and machine learning. MIT Press, Cambridge, Mass. ;, 2006.  ISBN 026218253X.
  * Rendell et al. [2011] L. Rendell, L. Fogarty, W. J. E. Hoppitt, T. J. H. Morgan, M. M. Webster, and K. N. Laland.  Cognitive culture: theoretical and empirical insights into social learning strategies.  _Trends in Cognitive Sciences_ , 15:68–76, 2011.
  * Rhodes and Baron [2019] M. Rhodes and A. S. Baron.  The development of social categorization.  _Annual review of developmental psychology_ , 1:359–386, 2019.
  * Sutton and Barto [1998] R. S. Sutton and A. G. Barto.  _Reinforcement Learning: An Introduction_.  Adaptive Computation and Machine Learning series. MIT Press, Cambridge, 1998.
  * Tolman [1948] E. C. Tolman.  Cognitive maps in rats and men.  _Psychological review_ , 55 4:189–208, 1948.  URL https://api.semanticscholar.org/CorpusID:42496633.
  * Wilson et al. [2014] R. C. Wilson, Y. K. Takahashi, G. Schoenbaum, and Y. Niv.  Orbitofrontal cortex as a cognitive map of task space.  _Neuron_ , 81:267–279, 2014.  URL https://api.semanticscholar.org/CorpusID:2713112.
  * Witt et al. [2023] A. Witt, W. Toyokawa, W. Gaissmaier, K. N. Lala, and C. M. Wu.  Social learning with a grain of salt.  In _Proceedings of the Annual Meeting of the Cognitive Science Society_ , volume 45, 2023.
  * Wu et al. [2018] C. M. Wu, E. Schulz, M. Speekenbrink, J. D. Nelson, and B. Meder.  Generalization guides human exploration in vast decision spaces.  _Nature Human Behaviour_ , 2:915–924, 2018.  URL https://api.semanticscholar.org/CorpusID:53278852.

##  Appendix A Representation entropy

Here we provide some additional details on our use of representation entropy
as a measure of cognitive cost. For a comprehensive introduction to entropy
(and other related information-theoretic quantities) we direct the reader to
Cover and Thomas [2006]. For a discrete random variable $X\in\mathcal{X}$ with
probability mass function $p(x):=\text{Pr}(X=x)$, the entropy of $X$ is given
by

$H_{b}(X)=-\sum_{x\in\mathcal{X}}p(x)\log_{b}(p(x))$ (7)

When $b=2$ (as it typically is), the entropy has units of bits. For a
continuous random variable $Y\in\mathcal{Y}$ with probability density function
$p(y)$, the differential entropy of $Y$ is given by

$H_{b}(Y)=-\int_{y\in\mathcal{Y}}p(y)\log_{b}(p(y))$ (8)

If $Y$ follows a multivariate Gaussian distribution with covariance $\Sigma$,
then $H(Y)$ is computed as

$H_{2}(Y)=\frac{1}{2}\log_{2}|\Sigma|+\frac{n}{2}\big{(}\log_{2}(2\pi
e)\big{)}$ (9)

where $n$ is the dimensionality of $Y$ [Rasmussen and Williams, 2006]. This
allows us to compute the entropy of our agent value functions $\mathbf{v}$,
since each value function is in effect a continuous random variable
distributed according to a multivariate Gaussian with known $\Sigma$. The
total entropy of a representation $\chi_{\text{ind}}$ that consists of the
value functions of $M$ agents over $n$ states is thus given by

$H_{2}[\chi_{\text{ind}}]=\frac{M}{2}\big{[}\log_{2}|\Sigma|+n\big{(}\log_{2}(2\pi
e)\big{)}\big{]}$ (10)

To compute the entropy of a state-aggregated value function estimate
$\hat{\mathbf{v}}$, we can treat $\hat{\mathbf{v}}$ as a new continuous RV
with a lower-dimensional multivariate Gaussian distribution whose covariance
$\Sigma_{\text{agg}}$ is determined entirely by $\Sigma$ and the level of
state aggregation. Determining $\Sigma_{\text{agg}}$ is then sufficient to
compute $H(\hat{\mathbf{v}})$. For the group-only representation cost, we need
simply to compute the entropy of the scalar discrete random variable
$\mathbf{z}$ (the agents’ group assignments). With a uniform distribution over
$K$ groups, this is simply given (following Eq. 7) by

$H_{2}(\mathbf{z})=-\sum_{k=1}^{K}\frac{1}{K}\log_{2}\big{(}\frac{1}{K}\big{)}=\log_{2}(K)$
(11)

The total entropy of a groups-only representation $\chi_{\text{grp}}$ is given
by combining this with the total entropy of the group mean value functions:

$H_{2}[\chi_{\text{grp}}]=\frac{K}{2}\big{[}\log_{2}|\Sigma|+n\big{(}\log_{2}(2\pi
e)\big{)}\big{]}+M\log_{2}(K)$ (12)


# Phone-to-audio alignment without text: A Semi-supervised Approach

###### Abstract

The task of phone-to-audio alignment has many applications in speech research.
Here we introduce two Wav2Vec2-based models for both text-dependent and text-
independent phone-to-audio alignment. The proposed Wav2Vec2-FS, a semi-
supervised model, directly learns phone-to-audio alignment through contrastive
learning and a forward sum loss, and can be coupled with a pretrained phone
recognizer to achieve text-independent alignment. The other model,
Wav2Vec2-FC, is a frame classification model trained on forced aligned labels
that can both perform forced alignment and text-independent segmentation.
Evaluation results suggest that both proposed methods, even when
transcriptions are not available, generate highly close results to existing
forced alignment tools. Our work presents a neural pipeline of fully automated
phone-to-audio alignment. Code and pretrained models are available at
https://github.com/lingjzhu/charsiu.

Index Terms—  forced alignment, phone segmentation, deep learning, automatic
speech recognition

##  1 Introduction

Aligning phones to audio has been a fundamental task in speech and phonetic
research, as many speech analyses depend on knowing the exact timing of
phones. There are currently a few existing tools for performing forced
alignment over phone sequences [1, 2, 3, 4, 5]. The majority of them are based
on the classic HMM system built upon Kaldi [6] or the HTK toolkit [7]. Though
neural ASR systems have shown better performance than HMM systems [8], phone-
to-audio alignment has not benefited a lot from the predictive power of neural
networks. One reason is that model ASR has increasingly shifted towards end-
to-end training using loss functions like CTC [9] that disregards precise
frame alignment. Only a few works explored using neural networks to perform
segmentation of sentences [10] and phones [11, 12, 13]. These works
demonstrate great potentials for neural forced alignment, but they still
required text transcriptions. Compared to forced alignment, text-independent
alignment has received much less attention. While a few unsupervised or self-
supervised phone segmentation methods have been proposed [14, 15, 16, 13],
most of them (except [13]) only output phone boundaries but do not jointly
predict both boundaries and phones, making them less practical in textless
alignment. Therefore, we also developed a method to perform alignments
directly on audio, as transcriptions might not always be available in many
practical applications.

In this study, we present two neural networks for both text-dependent and
text-independent phone-to-audio alignment. Our contribution can be summarized
as follows. First, we present a semi-supervised model to perform forced
alignment, but can also be combined with a phone recognizer for text-
independent alignment. It achieves comparable results with four existing
forced alignment tools. Secondly, we have also tested a frame classification
model for text-independent alignment, which achieved close performance to
text-dependent models. Extensive analyses show that the proposed methods still
maintain good performance in different settings. In contrast to currently
available tools, our work provides speech and language researchers with a
strong deep learning based tool to obtain phone segmentations even when
textual transcriptions are not available.

##  2 Neural forced alignment

Given a speech signal $X^{S}\in\mathbb{R}^{1\times T_{raw}}$ and a sequence of
phonetic symbols $Y^{P}\in\mathbb{R}^{1\times N}$, some neural networks are
used to encode them into hidden representations $\Phi$, such that
$\boldsymbol{Y}=\Phi_{S}(Y_{P})$ and $\boldsymbol{X}=\Phi_{P}(X_{S})$, where
$\boldsymbol{X}\in\mathbb{R}^{K\times
T}=[\boldsymbol{x}_{1},\boldsymbol{x}_{2},\dots,\boldsymbol{x}_{t}]$ and
$\boldsymbol{Y}\in\mathbb{R}^{K\times
N}=[\boldsymbol{y}_{1},\boldsymbol{y}_{2},\dots,\boldsymbol{y}_{n}]$, $T\geq
N$. We seek to learn a transformation $f$ to encode these features into the
monotonic alignment $S=f(\boldsymbol{X},\boldsymbol{Y})$, so that it aligns
each phonetic symbol $y^{P}_{m}$ to a subset of consecutive acoustic frames
$\boldsymbol{x}_{i}$: $S$ =
$[\boldsymbol{x}_{i}=y^{P}_{m},\boldsymbol{x}_{j}=y^{P}_{m},\dots,\boldsymbol{x}_{k}=y^{P}_{n}]$.
If transcriptions are not available, phone labels $\hat{Y}^{P}$ can be
estimated with a phone recognizer from audio.

The proposed neural forced alignment model learns the phone-to-audio alignment
through the self-supervised task of reconstructing the quantized embeddings of
original speech with both heavily masked speech representations and phonemic
information [12]. This could be implemented as the same pretraining task of
Wav2Vec2 [8]. Meanwhile, we relied on the forward sum loss to guide the
attention matrix to be monotonic alignments [17, 18].

###  2.1 Encoders

We used the pretrained wav2vec2-base model [8] with a convolutional layer on
top as the speech encoder $\Phi_{S}$. As the original frame rate of Wav2Vec2
(about 20ms per frame) could be too coarse to derive precise timing, we
increased the frame rate by two methods, either by reducing the stride of the
last convolutional layer in its feature extractor from 2 to 1 (denoted as
-10ms), or by upsampling the raw speech signal to 32kHz (denoted as -32k).
This effectively doubled the frame rate to 98Hz (about 10ms per frame), while
allowing the model to load all the pretrained parameters of wav2vec2-base.

The phone encoder $\Phi_{T}$ is a reduced version of the BERT model [19] with
a convolutional layer on top. It has 4 hidden layers, each of which has 12
self-attention heads of 32 hidden dimensions. The model was pretrained with
masked language modeling [19] on phone sequences generated from the whole Book
Corpus [20], which amounts to more than 2 billion of phone tokens.

###  2.2 Loss function

The model learns a matrix $\boldsymbol{A}\in\mathbb{R}^{N\times T}$ that
aligns the phone sequence $\boldsymbol{Y}$ and the masked speech
representation $\boldsymbol{\hat{X}}$.

$\boldsymbol{D}_{ij}=f_{y}(\boldsymbol{y}_{i})^{T}f_{x}(\boldsymbol{\hat{x}}_{j})$
(1) $\boldsymbol{A}=softmax(\boldsymbol{D},\texttt{dim=0})$ (2)
$\boldsymbol{H}=concatenate[\boldsymbol{\hat{X}},\boldsymbol{Y}\boldsymbol{A}]$
(3)

where $f_{y}$ and $f_{x}$ are two dense layers.
$\boldsymbol{D}\in\mathbb{R}^{N\times T}$ is the (unnormalized) similarities
between $\boldsymbol{Y}$ and $\boldsymbol{X}$. For the hidden states
$\boldsymbol{H}\in\mathbb{R}^{2K\times
T}=[\boldsymbol{h}_{1},\dots,\boldsymbol{h}_{T}]$ at each time step, the model
maximizes the similarity between each $\boldsymbol{h}_{t}$ and the quantized
version of the acoustic feature at the same time step. The quantized
embeddings can be selected from a codebook using Gumbel softmax parameterized
by a neural network, as in the original Wav2Vec2 [8]. Let $\boldsymbol{Q}_{t}$
represents the set consisting of a true quantized embedding and $n=50$
negative samples at step t. The loss function is formulated as follows.

$\mathcal{L}_{m}=\frac{1}{T}\sum_{t=1}^{T}-\log\frac{\exp{(sim(\boldsymbol{h}_{t},\boldsymbol{q}_{t})/\kappa)}}{\sum_{\boldsymbol{q}_{k}\sim\boldsymbol{Q}_{t}}\exp(sim(\boldsymbol{h}_{t},\boldsymbol{q}_{k})/\kappa)}$
(4)

The masking is performed with spectral augmentation with both time and feature
masking [21]. If the time masking probability is too low, the model might
simply ignore the phone information. However, if the time masking probability
is too high, it will create a gap between training and inference inputs, as
the input audio for inference is usually not corrupted. We sampled the time-
masking probability $p/100$ for each batch from a discrete uniform
distribution of $[p_{l},p_{h}]$. The original weights of the Wav2Vec2
quantizer were kept fixed throughout training.

Yet the attention matrices might not necessarily align acoustic frames and
phone tokens. Therefore we added the forward-sum loss used in HMM systems to
constrain the attention matrix to be monotonic and diagonal [17, 18].

$\mathcal{L}_{FS}=-\sum_{\boldsymbol{X},\boldsymbol{Y}\in S}\log
P(\boldsymbol{Y}|\boldsymbol{X})$ (5)

where $S$ is the optimal alignment. All probability quantities are estimated
in the attention matrix $\boldsymbol{A}$. This loss function can be
implemented using the off-the-shelf PyTorch CTC loss by setting the
probability of the blank symbol to an impossibly low value. We adapted the
implementation from [17, 18]. So the final loss is a weighted sum of the two
loss functions $\mathcal{L}=\mathcal{L}_{m}+\lambda\mathcal{L}_{FS}$, where
$\lambda$ was set to 1.

The alignment is notoriously hard to train, since early deviations from the
ground truth only get reinforced in subsequent iterations [22]. Techniques
such as guided attention loss [23] or diagonal prior [17, 18] or external
supervision with available alignments [24] have been used to enable the
formation of text-to-speech alignment. To facilitate alignment learning, we
train the model incrementally in the curriculum learning paradigm [25] on
increasingly longer audio samples. Like noisy student training [26] and the
iterative training in Kaldi [6], a series of models are generated in the
process.

  1. 1.

Initialize model $M$ with pretrained weights $M_{0}$. Set $M=M_{0}$.

  2. 2.

Partition the length-sorted data $C$ into $k$ different chunks
$C_{1},C_{2},\dots,C_{k}$ by duration, such that each chunk contains audio
longer than the previous chunk. Then train $M$ on each chunk sequentially with
a frame shift of 20ms.

  3. 3.

Upsample the frame shift of $M$ to 10ms. Retrain $M$ using the same data and
procedure as in 2.

  4. 4.

(Optional) Train a frame classification model $M^{\prime}$ from scratch with
alignments generated from the last step. Set $M=M^{\prime}$ and go back to 2.

Reinitialization with a frame classification model speeds up the convergence
as it has been noted that a prior is essential for obtaining good alignment
[22]. As the learned attention alignments are sharp, the decoding can be
performed simply using the argmax function without resorting to the Viterbi
decoding. This forced alignment is referred to as W2V2-FS.

##  3 Frame classification model

For text-independent alignment, we fine-tuned a Wav2Vec2 model to perform
frame-wise phone classification (referred to as W2V2-FC-Libris). The training
labels were alignments of the 960h Librispeech [27] obtained through MFA used
in [28]. The alignments were discretized into 10ms or 20ms frames and the
model was trained to classify each frame into a phone category using the
cross-entropy loss. During inference, the phone alignment was acquired through
model predictions directly with audio, but forced alignment was also performed
using a Dynamic Time Warping (DTW) algorithm based on an output probability
matrix and its transcription. We also trained additional frame classification
models (W2V2-FC) on the alignments generated from the semi-supervised method.

##  4 Experiments

Table 1: Evaluation results of text-dependent alignment Model | P | R | F1 | R-val | Overlap
---|---|---|---|---|---
FAVE | 0.57 | 0.59 | 0.58 | 0.64 | 74.3%
MFA-Libris | 0.61 | 0.61 | 0.61 | 0.67 | 73.5%
MFA | 0.62 | 0.63 | 0.63 | 0.68 | 75.0%
Gentle | 0.49 | 0.46 | 0.48 | 0.56 | 67.7%
WebMAUS | 0.70 | 0.70 | 0.70 | 0.75 | 78.8%
W2V2-FC-20ms-Libris | 0.49 | 0.47 | 0.48 | 0.56 | 73.8%
W2V2-FC-10ms-Libris | 0.57 | 0.54 | 0.55 | 0.62 | 76.4%
W2V2-FC-32k-Libris | 0.66 | 0.63 | 0.64 | 0.69 | 79.3%
W2V2-FS-20ms | 0.47 | 0.49 | 0.48 | 0.55 | 71.6%
W2V2-FS-10ms | 0.68 | 0.68 | 0.68 | 0.73 |  80.4%
W2V2-FS-32k | 0.63 | 0.65 | 0.64 | 0.69 | 79.3%
Pretrained G2P converter
W2V2-FS-20ms | 0.40 | 0.42 | 0.41 | 0.49 | 65.1%
W2V2-FS-10ms | 0.56 | 0.58 | 0.57 | 0.63 | 72.5
W2V2-FC-32k-Libris | 0.58 | 0.57 | 0.58 | 0.64 | 73.0%
Phone set adaptation (TIMIT-61)
W2V2-FS-20ms | 0.49 | 0.53 | 0.51 | 0.57 | 70.5%
W2V2-FS-10ms | 0.66 | 0.70 | 0.68 | 0.72 | 79.7%

For the English dataset, we used the Librispeech [27] the 1600h Common Voice
6.1 [29]. The TIMIT dataset with human annotations [30] was used for
evaluation. The original TIMIT 61 phone annotations were collapsed into the 39
CMU phone set. The flap [R] DX did not have a one-to-one mapping, so we kept
all occurrences. While this slightly degraded the overall accuracy, this was
the same for all forced alignment methods under comparison except WebMAUS,
which used more phones than those in the CMU phone set. All data and their
train/test partitions were based on the HuggingFace datasets package [31]. All
textual transcriptions were converted into phone sequences using an open
source grapheme-to-phoneme (G2P) converter [32].

In our experiment, we ran four iterations of training. We set the time masking
probability to $[0.05,0.2]$ for the 20ms models and $[0.05,0.4]$ for 10ms
models. The first two iterations were trained on Common Voice. We partitioned
the Common Voice dataset into three sets $\\{C1,C2,C3\\}$, $C1$ containing
sentences below 3 seconds, $C2$ with sentences between 3-5 seconds and $C3$
consisting of sentences up to 10 seconds. The third iteration was trained on
Librispeech and in the fourth iteration, the model was fine-tuned on the TIMIT
train set. The model generally converged within a thousand training steps so
training on full dataset was not needed. Annotated boundary information was
never used during training. Evaluation was performed on the TIMIT test set,
which does not overlap with any of the training data. The proposed model was
coded in pytorch based on the Wav2Vec2 implementations in the transformers
package. By default, we used an effective batch size of 32 and the Adam
optimizer with a learning rate of 1e-6 and rate decay of 1e-6. The training
was done distributedly on three GPUs with 12 GB of memory.

We compared our method with four publicly available forced aligners: Montreal
Forced Aligner (MFA) [1], Penn Forced Aligner (FAVE) [3], WebMAUS [4] and
Gentle [5]. Specifically for MFA, we tested its performance both using a
pretrained Librispeech model and performing forced alignment from scratch on
TIMIT.

##  5 Evaluation

Table 2: Evaluation results of text-independent alignment Model | P | R | F1 | R-val | Overlap
---|---|---|---|---|---
W2V2-CTC-10ms | 0.31 | 0.29 | 0.30 | 0.42 | 43.9%
W2V2-CTC-20ms | 0.31 | 0.30 | 0.31 | 0.42 | 46.6%
Phone recognition \+ W2V2-FS
W2V2-FS-20ms | 0.40 | 0.42 | 0.41 | 0.48 | 64.2%
W2V2-FS-10ms | 0.56 | 0.58 | 0.57 | 0.63 | 71.5%
W2V2-FC-32k-Libris | 0.57 | 0.57 | 0.57 | 0.64 | 72.2%
Direct inference |  |  |  |  |
W2V2-FC-20ms-Libris | 0.57 | 0.59 | 0.58 | 0.63 | 72.7%
W2V2-FC-10ms-Libris | 0.55 | 0.58 | 0.56 | 0.62 | 72.5%
W2V2-FC-32k-Libris | 0.60 | 0.63 | 0.61 | 0.66 | 74.3%

We evaluated the forced alignment results with precision, recall, F1, and
R-value [14]. For each predicted phone boundary, if the timing was within
tolerance $\tau$ and the predicted phone matches, it was considered a hit. As
each boundary marked the onset and the offset of consecutive phones, we only
evaluated the phone onsets with a tolerance of 20ms. In addition, we also
measured the percentage of correctly predicted frame labels at a 10ms time
scale.

In Table 1, all methods can perform forced alignment with decent accuracy,
though Gentle seems to lag slightly behind other methods. The proposed forced
alignment method also shows comparable performance, comparable to all existing
forced alignment methods, if not better. In general, increasing the frameshift
from 20ms to 10ms improved the alignment accuracy. However, phones in
continuous speech often undergo a variety of phonetic processes such as
deletion and insertion, resulting in pronunciation variants deviating from
dictionary pronunciations. Yet most G2P converters can only provide dictionary
pronunciations, which may lead to decreased performance. To shed light on this
issue, we tested the influence of a pretrained G2P converter. In Table 1, it
is clear that using the dictionary pronunciations decreases the alignment
performance by a significant margin, yet the performance is still close to
current forced aligners.

As different speech varieties may have varying phone inventories (like
American English vs. Received Pronunciation), it would be desirable for the
pretrained W2V2-FS model to adapt to different phone inventories. To test
this, we can simply fine-tune the model after modifying the phone embeddings
of the text encoder. For TIMIT, we expanded the existing 39 phone embeddings
of the original W2V2-FS model to include the full TIMIT 61 phone set while
keeping the rest of the pretrained weights intact, and fine-tuned the W2V2-FS
model. The model converged quickly after a few hundred iterations and achieved
comparable performance as the original model (see last two rows in Table 1).
Yet W2V2-FC models might be less flexible in this regard, as it requires pre-
existing alignments to train.

For text-independent alignment, we trained a phone recognizer on both
Librispeech and Common Voice by fine-tuning wav2vec2-base using the CTC loss.
This phone recognizer was used to derive phone transcriptions first before
performing forced alignment with W2V2-FS. Using a phone recognizer works
reasonably well (Table 2), as the alignment is only 1% worse than using a
pretrained G2P converter. As expected, the alignments generated by the CTC
loss alone deviated significantly from the ground truth, since the CTC loss
does not encourage time-synchronous alignment. In contrast, W2V2-FC models
achieved better performance (Table 2) than other text independent models,
highly comparable to that of text-dependent alignment methods. This proves
that the frame classification model used in HMM systems are still robust with
neural networks. Overall, these encouraging results suggest that both text
independent models are practical means to derive phone-to-audio alignments.

Fig. 1: Sample alignments. From top to bottom: 1) No curriculum training; 2) No $\mathcal{L}_{FS}$; 3) No $\mathcal{L}_{m}$; 4) Full model. Table 3: Evaluation of W2V2 models from different iterations Model | Iteration | Training Data | R-val | Overlap
---|---|---|---|---
FS-10ms | 1 | Common Voice | 0.41 | 60.7%
FC-10ms(w/o text) | 1 | Common Voice | 0.37 | 55.8%
FS-10ms | 2 | Common Voice | 0.51 | 68.1%
FC-10ms(w/o text) | 2 | Common Voice | 0.53 | 66.2%
FS-10ms | 3 | Librispeech | 0.60 | 74.1%
FC-10ms(w/o text) | 3 | Librispeech | 0.52 | 68.1%
FC-10ms(w/ text) | 3 | Librispeech | 0.63 | 75.3%

When the errors of G2P are taken into account, our forced alignment methods
might not possess a strong advantage over existing tools. The discrepancy
could be accounted for by two causes. First, silence poses challenges to
forced alignment as it is not encoded in the transcription. Without
preprocessing silence, phones are often misaligned when a long silent interval
is present. Secondly, we did not explicitly model pronunciation variants.
However, HMM-based systems [6] often have built-in methods to alleviate these
two issues. This will be the next step for our research.

##  6 Ablation analysis

We performed ablation analysis to examine the effectiveness of the W2V2-FS
model. As shown in Fig 1, both the contrastive loss and the forward-sum loss
are necessary to facilitate the learning of sharp attention alignment. If the
audio and the transcription were aligned using only the forward sum loss, the
learned alignment was noisy. If either the curriculum learning or the
contrastive loss is removed, the model generally fails to converge to a
diagonal alignment pattern.

Multiple iterations of training are generally needed to boost performance (see
Table 3), as initial alignments tend to be suboptimal, especially for noisy
datasets like Common Voice. During training, we found that the forced
alignment model was very sensitive to initialization, and poor alignment was
reinforced but not corrected in subsequent training. So using a frame
classification to provide a good prior estimation was important. Audio quality
had a large impact on the quality of alignment. Training on TIMIT or
Librispeech-clean from scratch ended up with better alignments, faster
convergence rates, and fewer iterations than training on the crowd-sourced
Common Voice. Table 3 illustrates the variations of model performance when
noisy data and audio domain mismatch are both present, though the performance
gaps can be narrowed through fine-tuning.

##  7 Discussions and Conclusions

In this study, we present two deep learning-based methods to align audio to
phones and achieve comparable performance against several existing forced
alignment tools. These models can be combined as a pipeline to bootstrap phone
labels from naturalistic audio data in the wild, which are massive but remain
under-exploited. This will have implications for speech corpus creation,
phonetic research, and interdisciplinary studies. Given the good performance
of W2V2-FC-Libris, our method can also be combined with MFA[1] to train a
large-scale text-independent frame classification model to get rid of text in
inference. The W2V2 models can be accelerated by harnessing the computing
power of GPU. For example, in the inference mode, our model aligned the 1600
hours Common Voice dataset in about 10 hours on a 2080Ti GPU without batching.
While only tested on phone segmentation, the model also performed well on
related tasks such as word or sentence segmentation. In the future, we will
continue to improve the alignment accuracy and extend it to different
languages.

## References

  * [1] M. McAuliffe et al.,  ``Montreal forced aligner: Trainable text-speech alignment using kaldi.,''  in INTERSPEECH, 2017, vol. 2017, pp. 498–502.

  * [2] K. Gorman et al.,  ``Prosodylab-aligner: A tool for forced alignment of laboratory speech,''  Canadian Acoustics, 2011.
  * [3] I. Rosenfelder et al.,  ``FAVE (forced alignment and vowel extraction) program suite,''  URL http://fave. ling. upenn. edu, 2011.
  * [4] T. Kisler et al.,  ``Signal processing via web services: the use case webmaus,''  in Digital Humanities Conference 2012, 2012.
  * [5] R. Ochshorn and M. Hawkins,  ``Gentle,'' https://lowerquality.com/gentle/, 2017.
  * [6] D. Povey et al.,  ``The kaldi speech recognition toolkit,''  in ASRU, 2011.
  * [7] S. Young et al.,  ``The HTK book,''  Cambridge university engineering department, vol. 3, no. 175, pp. 12, 2002.
  * [8] A. Baevski et al.,  ``wav2vec 2.0: A framework for self-supervised learning of speech representations,''  in NeurIPS, 2020.
  * [9] A. Graves et al.,  ``Connectionist temporal classification: labelling unsegmented sequence data with recurrent neural networks,''  in ICML, 2006.
  * [10] L. Kürzinger et al.,  ``CTC-segmentation of large corpora for german end-to-end speech recognition,''  in International Conference on Speech and Computer, 2020.
  * [11] M. Kelley and B. Tucker,  ``A comparison of input types to a deep neural network-based forced aligner,''  in INTERSPEECH, 2018.
  * [12] K. Schulze-Forster et al.,  ``Joint phoneme alignment and text-informed speech separation on highly corrupted speech,''  in ICASSP, 2020.
  * [13] Y. Teytaut and A. Roebel,  ``Phoneme-to-Audio Alignment with Recurrent Neural Networks for Speaking and Singing Voice,''  in INTERSPEECH, 2021.
  * [14] F. Kreuk et al.,  ``Self-Supervised Contrastive Learning for Unsupervised Phoneme Segmentation,''  in INTERSPEECH, 2020.
  * [15] H. Kamper and B. van Niekerk,  ``Towards unsupervised phone and word segmentation using self-supervised vector-quantized neural networks,''  arXiv preprint arXiv:2012.07551, 2020.
  * [16] S. Bhati et al.,  ``Segmental contrastive predictive coding for unsupervised word segmentation,''  arXiv preprint arXiv:2106.02170, 2021.
  * [17] K. Shih et al.,  ``Rad-tts: Parallel flow-based tts with robust alignment learning and diverse synthesis,''  in ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, 2021.
  * [18] R. Badlani et al.,  ``One tts alignment to rule them all,''  arXiv preprint arXiv:2108.10447, 2021.
  * [19] J. Devlin et al,  ``BERT: Pre-training of deep bidirectional transformers for language understanding,''  in NAACL, 2019, pp. 4171–4186.
  * [20] Y. Zhu et al.,  ``Aligning books and movies: Towards story-like visual explanations by watching movies and reading books,''  in ICCV, 2015.
  * [21] D. Park et al.,  ``SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition,''  in INTERSPEECH, 2019.
  * [22] A. Zeyer et al.,  ``CTC in the Context of Generalized Full-Sum HMM Training,''  in INTERSPEECH, 2017.
  * [23] H. Tachibana,  ``Efficiently trainable text-to-speech system based on deep convolutional networks with guided attention,''  in ICASSP, 2018.
  * [24] Y. Ren et al.,  ``Fastspeech 2: Fast and high-quality end-to-end text to speech,''  in ICLR, 2020.
  * [25] Y. Bengio et al.,  ``Curriculum learning,''  in ICML, 2009.
  * [26] D Park et al.,  ``Improved noisy student training for automatic speech recognition,''  in INTERSPEECH, 2020.
  * [27] V. Panayotov et al.,  ``Librispeech: an asr corpus based on public domain audio books,''  in ICASSP, 2015.
  * [28] L. Lugosch et al.,  ``Speech Model Pre-Training for End-to-End Spoken Language Understanding,''  in ITERSPEECH, 2019.
  * [29] R. Ardila et al.,  ``Common voice: A massively-multilingual speech corpus,''  in LREC, 2020.
  * [30] J. Garofolo et al.,  ``DARPA TIMIT acoustic-phonetic continous speech corpus CD-ROM. NIST speech disc 1-1.1,''  NASA STI/Recon technical report, 1993.
  * [31] Q Lhoest et al.,  ``Datasets: A community library for natural language processing,''  arXiv preprint arXiv:2109.02846, 2021.
  * [32] Kyubyong Park and Jongseok Kim,  ``g2pE,'' https://github.com/Kyubyong/g2p, 2019.