There are many advancements in new model architectures in AI domain. Let me overview these advancements in this post.
Linear Compression Embedding
LCE [1] is simply using a matrix to project one embedding matrix to another: 
, where 
.
Pyramid networks, inception network, dhen, lce
Perceiver and Perceiver IO
Perceiver-based architectures [5,6] solve a pain point of traditional transformers: quadratic computation in terms of sequence length for computing KQV attention, where K/Q/V all have the shape batch_size x seq_len x emb_dim. Perceiver-based architectures creates a learnable V matrix which is independent of input sequence lengths, thus to reduce the attention computation to be linear. Thanks to linear attention computation time, within the same computation budget, it is also able to stack more attention blocks to improve the performance. There are three web tutorials that are helpful: [3,4,5].
Since Perceiver has done a good job of reducing attention computation to linear time (in terms of sequence length), I was wonder why recent LLMs did not, if any, adopt this architecture. A reddit online discussion shed some light [7].
Mamba
Mamba [9] is a new sequence modeling architecture which also achieves linear time based on so called “selective structured state space model”. Currently, it achieves SOTA in many tasks. The official implementation can be found here. There is not much material which can clearly articulate the theory behind Mamba. I have only found pieces of information here and there: [8, 10]. The way I understand about Mamba after watching [8] is:
- Mamba enjoys the best of both worlds of RNN and Transformers
Training Inference RNN O(N) but can’t be parallelized because states need to roll out sequentially O(1) per token Transformer O(N^2) but can be parallelized O(N) per token (see time analysis in [12]) Mamba O(N) and can be parallelized O(1) per token - According to basic definition of structured state space sequence (S4) models [10], Mamba has the state transition represented as differential equations:
Note,
is the change (i.e., derivative) of the hidden state. - The differential equations describe a real-time system but in sequence modeling we are interested in sequences with discrete time steps (i.e., tokens). “Discretization” in the paper means that we use approximation methods to derive new matrices
and 
when we assume the input function 
becomes the sequence 
, where 
is the step size.
There are two methods to do the discretization. [8] introduces the Euler’s method, which is simple for exposition. The paper uses zero-order hold, which is explained in more details in [13]. We use the Euler’s method for illustration, following [8]:- by definition of derivative we know that
- Substitute the state space model’s equations with step 1:
- by definition of derivative we know that
- State transitions during training can be parallelized without needing to wait for sequential rollout because the state transition model can be converted into convolutional formulation that only depends on the input sequences:
- A S4 mdoel is used for one input dimension. So if a token has 512 embedding size, there will be 512 S4 models in parallel. In comparison, transformers use attention heads where each attention head consumes all dimensions (or a group of dimensions) of a token.
Reference
- LCE explanation: https://fb.workplace.com/groups/watchrankingexperiments/permalink/2853044271454166/ (internal) and https://www.dropbox.com/scl/fi/0ygejam1mg2jhocfmgoba/watch-_-chaining-ranking-experiment-review-_-group-_-workplace.pdf?rlkey=4euhhg6f0h838ythg1na1amz8&dl=0 (external)
- https://medium.com/@curttigges/the-annotated-perceiver-74752113eefb
- https://medium.com/ml-summaries/perceiver-paper-s-summary-3c589ec74238
- https://medium.com/ml-summaries/perceiver-io-paper-summary-e8f28e451d21
- Perceiver: General Perception with Iterative Attention: https://arxiv.org/abs/2103.03206
- Perceiver IO: A General Architecture for Structured Inputs & Outputs: https://arxiv.org/abs/2107.14795
- https://www.reddit.com/r/MachineLearning/comments/tx7e34/d_why_arent_new_llms_using_the_perceiver/
- https://www.youtube.com/watch?v=8Q_tqwpTpVU
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces: https://arxiv.org/abs/2312.00752
- https://srush.github.io/annotated-s4/?ref=blog.oxen.ai
- https://gordicaleksa.medium.com/eli5-flash-attention-5c44017022ad
- Llama code anatomy: https://czxttkl.com/2024/01/17/llama-code-anatomy/
- https://en.wikipedia.org/wiki/Discretization#Derivation
. You would like to express the distribution of 
. You have some prior 
. You would like to know the distribution of the posterior distribution 
. Based on Bayes’ theorem,![\[p(\theta|x)=\frac{p(\theta) p(x|\theta)}{p(x)}=\frac{p(\theta) p(x|\theta)}{\int_\theta p(\theta) p(x|\theta) d\theta}\]](https://czxttkl.com/wp-content/ql-cache/quicklatex.com-2529e21a8f21d8b383b095d0962a740d_l3.png "Rendered by QuickLaTeX.com")
is often hard to compute. Therefore, it is hard to know the exact statistics of the posterior distribution, such as the mean of the posterior distribution. MCMC has some mechanism to allow you sample points on the space of 
while 
is a fixed variance. According to basic conjugate prior mathematics [4], you have a prior distribution 
which is also a Normal distribution. You have observed 3 students’ scores therefore you can express the likelihood function as: 
. In this simple example, there is an analytical expression for the posterior distribution 
which is also a Normal distribution. But for our pedagogical purpose, let’s assume the posterior distribution 
is impossible to compute and we have to use MCMC to infer posterior statistics such as the posterior mean. 
at first.
is defined in [1].
and 
. ![\[\alpha(\theta^{new}, \theta^{old})=min \left(1, \frac{p(\theta^{old}|x)g(\theta^{new}|\theta^{old})}{p(\theta^{new}|x) g(\theta^{old}|\theta^{new})}\right)\]](https://czxttkl.com/wp-content/ql-cache/quicklatex.com-75a1b765204a57aa4e742b7c38f731ef_l3.png "Rendered by QuickLaTeX.com")
for 
. Each behavior cloning model models the output action as a normal distribution with a predicted mean and standard deviation. Each time when we want to sample an action, we first randomly pick a behavior cloning model and then sample the action multiple times according to the predicted mean and standard deviation. However, there is some additional adaptation for how to sample from the predicted mean and standard deviation. First, the standard deviation is scaled larger so that action sampling can be more diverse and aggressive. Second, to counter sampling bad actions, we learn a Q-function from logging data and use it to finally select the action to proceed the trajectory simulation. Although the Q-function estimates the Q-value for the behavior policy instead of the policy being learned, the authors argue that it “provides a conservative but relatively reliable long-term prior information”. The two points are manifested in Eqn. 3:
to prune trajectories with overall uncertainty larger than that. The way to quantify uncertainty of a trajectory is to look at the sum of each step’s uncertainty defined as 
, the largest disagreement between any two dynamic models in the ensemble.
Note one detail at line 14: if we simulate 
steps, we can get 
, the maximum standard deviation of the learned models in the ensemble. As the authors argue, “empirically the learned variance often captures both aleatoric and epistemic uncertainty, even for learning deterministic functions (where only epistemic uncertainty exists). We use the maximum of the learned variance in the ensemble to be more conservative and robust.”
. We have a dataset 
which contains tuples of 
for 
. These are observed features and labels. If we apply a Bayesian logistic regression on such data, the logistic regression will have a weight vector of the same dimension as the features: 
. The prediction function is a sigmoid function: 
.
given the observation 
. We assume that 
. 
, what would be 
and 
?
achieves maximum 
, which is able to be found using gradient ascent or expressed in a closed formula (I think 
. Because 
, we can further simplify to 
. 
can actually be seen as the logarithm of another Gaussian distribution 
, whose mean is at 
, and 
absorbs any remaining constants. This is why Algorithm 3 updates 
because 
is essentially from 
since this is essentially 
.
, with 
being the number of total nodes and 
is the vector dimension. We hope to learn these vector representations using 
,
is the adjacency matrix, 
is the node representation from the previous layer, 
is an activation function, and 
is the parameters of this layer. 
is some basic feature representation of nodes, which can be as simple as an identity matrix, or as complex as shortest path distance to other nodes.
does not involve each node’s own representation; (2) nodes with large degrees will tend to have large values in their feature representation while nodes with small degrees will have small values. Therefore, a better propagation rule is:
,
, 
is the inverse degree matrix of 
as a way to normalize the following value by the degree of each node.
, which has some theoretical connection to “localized spectral filters”. This form of propagation rule, according to [5], not only takes into considerations the degree of the source node but also the target node for each node pair.
can be learned by unsupervised learning or semi-supervised learning manner. What we describe above is the unsupervised learning approach, in which node embeddings are computed without knowing any node labels. A more interesting point of [3] is that the propagation rule can be used to learn node representation in a semi-supervised learning setting, where only some nodes’ labels are revealed but the representations of all the nodes in the graphs can be learned. This is because if we back-propagate on the prediction loss function, it can eventually back-propagate to all other nodes due to that the nodes with labels have representations as the aggregation from other nodes’. 
is the distribution of user 
‘s played genres at time 
. Hence, ![\mathbf{n}^t_i=[n^t_{i1}, \cdots, n^t_{iN}]](https://czxttkl.com/wp-content/ql-cache/quicklatex.com-15ceafaa5104efb08d21a81fb48771d2_l3.png "Rendered by QuickLaTeX.com")
normalized by total play counts 
: 
. The transitions between genres at different time steps are considered as a graph. 
:
, 
, is sampled from a Poisson distribution determined by 
, will then be determined by both the total play count 
. 
transition matrix to be learned denoting how user genre preferences change in two consecutive steps. The paper chooses to use some two-step maximal likelihood method to learn 
.