Temporal Latent Bottleneck: Synthesis of Fast and Slow Processing Mechanisms in Sequence Learning
Abstract
Recurrent neural networks have a strong inductive bias towards learning temporally compressed representations, as the entire history of a sequence is representedby a single vector.
By contrast, Transformers have little inductive bias towardslearning temporally compressed representations, as they allow for attention overall previously computed elements in a sequence.
Having a more compressed representation of a sequence may be beneﬁcial for generalization, as a high-levelrepresentation may be more easily re-used and re-purposed and will contain fewerirrelevant details.
At the same time, excessive compression of representationscomes at the cost of expressiveness.
We propose a solution which divides computation into two streams.
A slow stream that is recurrent in nature aims to learn aspecialized and compressed representation, by forcing chunks of K time steps intoa single representation which is divided into multiple vectors.
At the same time,a fast stream is parameterized as a Transformer to process chunks consisting of
K time-steps conditioned on the information in the slow-stream.
In the proposedapproach we hope to gain the expressiveness of the Transformer, while encouraging better compression and structuring of representations in the slow stream.
Weshow the beneﬁts of the proposed method in terms of improved sample efﬁciencyand generalization performance as compared to various competitive baselines forvisual perception and sequential decision making tasks.
Introduction
The interplay between fast and slow mechanisms for information processing and perception has beenstudied in both cognitive science and machine learning [5, 34].
In the brain, short-term and long-termmemory have developed in a specialized way.
Short-term memory is allowed to change very quicklyto react to immediate sensory inputs and perception.
It also tends towards high capacity storage ofall pieces of information which may be relevant for future reasoning [41, 3, 4].
By contrast, longterm memory changes slowly [42, 40], is highly selective and involves repeated consolidation.
Itcontains a set of memories that summarize the entire past, only storing details about observationswhich are most relevant [23, 7].
Deep Learning has seen a variety of architectures for processing sequential data [35, 58, 16].
Sequence models which only rely on a short ﬁxed context window may be seen as a form of fastprocessing mechanisms, because there is no dependence on the longer-term history.
Early researchin deep learning found that robustness is substantially improved by using a recurrent neural networkto represent the entire history of the sequence, with direct evidence found in the domain of handwriting synthesis [26].
These recurrent neural networks which force the entire hidden state to becompressed into a single hidden state became highly successful across speech recognition, naturallanguage processing, and other sequence domains [27, 28, 29, 59].
Recurrent neural networks alsosaw signiﬁcant advances in architectural inductive biases, such as a bias towards making small and
[Figure]
[Figure]
selective changes to the hidden state on each step which helped improve its generalization performance [45, 31, 25].
Despite their wide applicability, recurrent neural networks have at the same time struggled in modeling long-sequences, especially in the presence of long-range dependencies. To help address this,attention was added to recurrent neural networks to allow for information processing which can selectively bypass the recurrent state [8].
In later works the recurrent state was removed entirely infavor of the Transformer architecture, which shares information between positions using attention[61]. Research on scaling laws has shown that the expressiveness advantage of Transformers overrecurrent neural networks grows with the length of the sequences being modeled [20, 55, 11].
Transformers have become the dominant architecture across a wide range of domains includingvision [21], natural language [20, 55, 11, 67, 17, 56], and reinforcement learning [13, 39].
Despitetheir success, it is well known that Transformers are extremely data hungry and work well mainlyat scale [21, 46, 49]. This may be attributed to the inductive bias in Transformers to model allpossible pairwise interactions between the inputs which also results in a very high computationalcomplexity that scales quadratically with the input size.
The lack of selectivity in interactions inthe Transformer makes them extremely memory inefﬁcient and may bias them towards modelingunnecessary, irrelevant, or noisy interactions. This motivates the use of a Temporal Latent Bottleneck , with the goal of improving Transformerswhile using lesser data and memory.
The proposed model aims at introducing selectivity in theinteractions by dividing computation into two streams - (1) A high-level slow stream consisting of aset of vectors updated in a recurrent manner (also referred to as the temporal latent bottleneck), and,(2)
A low-level attention-based fast stream that processes the input. The fast stream processes locallyneighbouring information within chunks, while the slow stream contains information about distanttokens across chunks.
The slow stream and the fast stream interact using a multi-head attentionmechanism to achieve selectivity in how local and distant information is mixed.
We show that theresulting model substantially outperforms Transformers and demonstrates improved generalization,especially when the test data includes novel challenges for the model that were not encounteredduring training.
Methodology
We now present the proposed approach in detail.
Our model jointly leverages the strengths of Transformers [61] and recurrent neural networks [16, 35].
The fast stream operates on raw inputs and isinstantiated using a Transformer.
We denote this Transformer by F .
The slow stream operates at ahigher level and updates at a lower frequency than the fast stream.
Any recurrent function can beused to instantiate the slow stream.
2.1 Desiderata for Fast and Slow Streams of Processing
We give the detailed description of the proposed model in the next section.
Here, we give anoverview of our architecture and discuss some of its key properties.
We present a detailed dia- gram of our model in Figure 1.
Given an input sequence, it is ﬁrst divided into chunks of size K .
Each chunk is processed by a Transformer (denoted as F ) which is also known as the perceptualmodule (since it processes the raw input) or the fast stream.
While processing each chunk, F isalso conditioned on information from the slow stream G which is also called the temporal latentbottleneck module.
The slow stream is a recurrent stream which has its own state consisting of a setof N vectors also called temporal latent bottleneck state denoted as I in Figure 1.
In the followingsections, we use the term temporal latent bottleneck to refer to the temporal latent bottleneck state I .
This state is updated once per chunk using information from the perceptual module through across attention mechanism.
The perceptual module operates within each chunk while the temporal latent bottleneck operatesacross chunks slowly updating itself after each chunk has been processed by the perceptual module.
Thus, the only way the perceptual module gets information about inputs beyond its own chunk isthrough the temporal latent bottleneck.
One advantage of this is that the computational complexityof the attention mechanism in the proposed model is O ( K 2 + KN ) while that of a Transformer is O ( T 2 ) , where T is the length of the sequence, K is the chunk size, and N is the number of temporallatent bottleneck state vectors.
Since K << T and N << T , we can see that K 2 + KN < T 2 .
Therefore the proposed model has a much lower computational complexity compared to a Transformer.
Furthermore, the capacity of the temporal latent bottleneck is limited and much smaller thanthat of the perceptual module.
This encourages the temporal latent bottleneck to represent the mostsalient information about the past while the perceptual module represents only local information.
This creates an information asymmetry between the two streams.
This information asymmetry leadsto the perceptual module having a ﬁne grained view of the nearby inputs but a very coarse grainedview of the distant past.
This is very different from the usual self-attention which attends to alltokens in the sequence at the same level of granularity.
One advantage of having a bottlenecked view of the past is that it allows the model to forget irrelevantinformation.
For example, if an agent is navigating in a large maze, it does not need to have ﬁnegrained knowledge of its actions from the distant past.
In the case of a Transformer, it would attendto every step from the past (including steps from the distant past) which may be irrelevant in thepresent context thus wasting its capacity in modeling irrelevant details.
Introducing specialized fast and slow processing mechanisms in machine learning is not completelynew.
A number of works [66, 51, 63, 65, 62, 64] build on vision Transformers [21] by introducingconvolution-like hierarchies in Transformers.
The main idea behind these models is to progressivelyreduce the size of the feature maps similar to convolutions so that some parts of the network canspecialize at a coarser granularity. Our way of specializing fast and slow processing is different. We do not progressively reduce the size of the feature maps in deeper layers, rather informationﬂows both ways between the fast processing stream and the slow processing stream. This allows thehigher level (i.e. the slow stream) to convey information about the long-term context to the lowerlevels (i.e. the fast stream) as top-down conditioning information .
However, previous works on introducing hierarchies in vision Transformer do not allow the higher levels to convey any informationto the lower levels. Past works have shown the usefulness of top-down information sharing betweenspecialized processing streams. [53] explored an architecture in which multiple recurrent streamsare active at every time-step, but encouraged one stream to be more active by providing it directaccess to the input.
This selective sharing of information between the streams using attention ledto better generalization. [22] showed that top-down sequential feedback in Transformers, althoughcomputationally expensive, improves the learning of long term dependencies. Through our experiments we show that the proposed model substantially outperforms Transformerson tasks where generalization to variations unseen during training is critical. We verify this improvement in imitation learning (behavior cloning), where the training and evaluation distributionscan diverge as a result of the agent visiting states which are not present in the dataset.
Next, wedescribe the detailed implementation of our model.
2.2 Computational Steps
We ﬁrst describe our notation. We denote the input X as a sequence of T tokens - X =[ x 0 , x 1 , x 2 , . . .
, x t ] .
We chunk this input into chunks of size K resulting in (cid:98) T / K (cid:99) chunks.
Werefer to l th chunk as X l .
We represent the state of the temporal latent bottleneck I (i.e. the slowstream) as a set of M d -dimensional vectors.
As mentioned previously, we denote the temporal latent bottleneck module as G and the perceptual module as F .
G updates the temporal latent bottleneck state while F processes chunks X l to form the latent representation ¯ X l -
[Equation]
Module
[Equation]
Next, we go into details of the implementation of these modules.
Preliminaries .
The central component of our model is the key value attention mechanism [8, 61]. The attention mechanism allows to dynamically select information from a set of N read vectors R to update a set of M write vectors W .
The write vectors are ﬁrst projected to a set of queries Q = W q W , where Q ∈ R M× D . The read vectors are projected into keys K = W k R and values V =
W v R , where K ∈ R N× D and V ∈ R N× D .
The updated write vectors ¯ W are obtained as -
[Equation]
The ﬁrst term QK T D is a dot product that decides the attention scores given to each query-key pair.
The softmax normalizes the scores across the key dimension.
Pairs with high similarity will beassigned a high score.
The ﬁnal update to the write vector is obtained by taking the convex combination of values weighted by the attention scores.
The convex combination ensures that the dynamicselection of information is differentiable and can be trained by backpropagation.
The updated vectors ¯ W are utilized by residually adding them to the write vectors ( W = W + ¯ W ). If the readand write vectors refer to the same set of vectors, the attention mechanism is called a self-attentionmechanism [61]. If the read and write vectors refer to seperate set of vectors, the attention mechanism is called a cross attention mechanism [25, 24, 38]. In this paper, we use both kinds of attentionmechanisms.
The Transformer architecture is another central component of our model which stacks a series ofself attention and FFN layers. Each FFN layer is a 2 layered MLP that projects the input to a largerdimension and projects it back to the input dimension. In this case, we use the gelu activation [32]for the FFN. [Equation]
Given the attention and FFN layers, a single Transformer layer can be written as -
[Equation]
where LN refers to layer normalization layers [6].
In general, multiple such Transformer layers arestacked to solve any task.
We augment the Transformer, which has only self attentions, to include cross attentions too so thatthe temporal latent bottleneck module and perceptual module can share information with each other.
We now describe the speciﬁcs of the perceptual module and the temporal latent bottleneck module.
Perceptual Module F .
As mentioned previously, the perceptual module refers to the fast streamthat acts directly on the input.
The perceptual module operates on each chunk separately.
Therefore,at any time the input to the perceptual module are the tokens corresponding to a particular chunk X l = [ x l × K , x l × K +1 , . . .
, x l × K + K ] .
The perceptual module is a Transformer with self attentionlayers, cross attention layers, and FFNs.
It has 2 kinds of layers - (1) SELF ATTENTION + FFN;(2) CROSS ATTENTION + FFN.
The SELF ATTENTION + FFN layers process the input tokens asdescribed in Equation 5 and the CROSS ATTENTION + FFN layers integrate top-down informationfrom the temporal latent bottleneck state
I as follows -
[Equation]
We include one CROSS ATTENTION + FFN layer per R SELF ATTENTION + FFN layers.
Thediagramatic representation of the perceptual module is presented in Figure 1 (in the processing ofchunk X l ).
In the ﬁgure, we set R = 1 .
Temporal Latent Bottleneck Module G .
The temporal latent bottleneck (TLB) module representsthe slow stream that operates on the temporal latent bottleneck state I .
I is updated using information from a particular chunk processed by the perceptual module.
This update happens once foreach chunk of the perceptual module resulting in (cid:98) T / K (cid:99) updates for I .
Since the temporal latentbottleneck state I updates at a lower frequency than the perceptual module, it is expected to capture more stable and slowly changing features while the perceptual module captures faster changingfeatures resulting in multiple scales of information representation.
An update to the temporal latent bottleneck state I consists of a cross attention operation where the queries come from I andthe keys and values come from the output of the perceptual module.
This cross attention operationis followed by an FFN update to I .
Consider the perceptual module outputs for a chunk l to be¯ X l = [¯ x l × K , . .
, ¯ x l × K + K ] .
The update operation is implemented as follows:
[Equation]
180 The temporal latent bottleneck module introduces the notion of recurrence in our model.
We showthe details of this module in Figure 1 (inside the circle).
Perceptual Module + Temporal Latent Bottleneck Model.
We now present our complete architecture integrating both the perceptual module and the temporal latent bottleneck together.
Given asequence of tokens X
=
[ x 0 , x 1 , x 2 , . .
, x t ] .
We chunk this input into chunks of size K resultingin (cid:98) T / K (cid:99) chunks.
The chunks are processed sequentially one after the other.
For a chunk k , itis ﬁrst processed using the perceptual module conditioned on information from the temporal latentbottleneck state.
The outputs of the chunk are used to update the temporal latent bottleneck state
I .
The resultant temporal latent bottleneck state is then used to process the next chunk.
The fullmodel is presented in Figure 1.
We use a Transformer as the perceptual module in our experiments.
Thus our main contribution is introducing a temporal latent bottleneck into Transformers and showing its advantages through a variety of experiments.
We also present the detailed algorithm for theproposed approach in the Appendix (Algorithm 1).  Related Work
Hierarchical or Multiscale Recurrent neural networks.
This work takes inspiration from awide array of work on introducing multiple scales of processing into recurrent neural networks[18, 33, 54, 57, 43].
These works divide the processing into multiple streams each operating at adifferent temporal granularity.
While these works mainly focus on recurrent neural networks andtheir application is mainly on natural language tasks, we focus on introducing multiple streams ofprocessing and a hierarchical structure into Transformers while also focusing on a broader range ofdomains beyond natural language. Transformers. Some of the components we describe in the proposed model have been used previously in various Transformer models. Transformer XL [19] also divides the input into segments.
Each segment considers the tokens from the current segment and the previous segmentfor attention without passing gradients into the previous segments.
A number of previous works[66, 51, 63, 65, 62, 64] have worked on introducing a hierarchical structure in Transformers mainlyin the domain of vision.
The main goal of these works has been to introduce convolution-like hierarchies into Vision Transformers [21].
While these works progressively reduce the spatial resolutionof the inputs in order to introduce hierarchies, we introduce hierarchies by adding another slowstream of information processing and without reducing the spatial resolution of the inputs.
We alsoprovision for the higher level of the hierarchy (i.e. the slow stream) to provide information to thelower levels as top-down conditioning which is not possible in any of the previous works.
Top-Down Conditioning.
Top-down information is information propagated from higher to lowerlevels of the network.
It represents the models beliefs of the world and provides context for interpreting perceptual information.
[53] and [22] have shown the advantages of top-down conditioningin recurrent neural networks and Transformers respectively.
These works focus on different streamsof processing operating at the same temporal granularity and the top-down conditioning is providedby higher level streams to the lower level streams.
In our case, the top-down conditioning for theperceptual module is provided by the high-level slow stream which operates at a slower temporalgranularity.
This allows the perceptual model to be affected by much more long term high levelinformation as compared to just short-term high level information as in the case of [53] and [22].
The proposed model is similar to a parallel work called block recurrent Transformers [36].
Thereare few differences between our work and theirs.
First, they use a sliding window attention, whilewe divide the input into chunks.
In their paper, they perform cross attention and self attention inparallel while we ﬁnd that doing them sequentially and performing cross attention once per R selfattention steps yields better results.
They also use special tricks to deal with some instabilities intheir case, while we ﬁnd no such instabilities in our model.
Also, while their focus is mainly onnatural language tasks, we focus on a broader variety of tasks.
[-We defer the rest of the discussionon related works to Appendix Section 
-]
Experiments
[Table]
In this section, we outline the tasks in which weapplied the temporal latent bottleneck and direct the reader to the appendix for more detailson the experiments.
Our goal is to show thewide applicability and beneﬁts offered by thetemporal latent bottleneck , which we refer to as TLB.
We demonstrate that the proposed modeloutperforms competitive baselines across manydomains including vision, reinforcement learning, and natural language.
Our main goal is toshow that the proposed approach has high expressive power like Transformers while also being sample efﬁcient unlike Transformers.
Thus ourmain baselines are based on the original Transformer architecture.
For example, we compare against ViT [21] in image classiﬁcation, Decision Transformer [13] in Reinforcement Learning, and Vanilla Transformer in rest of the tasks.
We also compare against some representative baseline that offersome of the key properties that our model offers.
For example, we compare against state-of-the art Swin Transformer [50] which is a strong baseline for image classiﬁcation and is also hierarchicalsimilar to the proposed model.
We also compare against Transformer LS [68] which also processeslong-term and short-term information using different attention streams.
Furthermore, we also compare against Transformer-XL, which also divides the input into segments similar to the proposedmodel.
Another key point of the proposed model is that any position cannot attend to any information from the future beyond its chunk since the temporal latent bottleneck only summarizes the past,not the future .
Meanwhile, all the baselines we consider have bidirectional context i.e. they canattend to all of the past and the future.
We observe that despite this limitation, the proposed modeloutperforms all the considered baselines.
4.1 Temporal Latent Bottleneck For Perception
[Table]
Image Classiﬁcation.
Recently, Transformers havebeen widely applied for visual perception and haveshown strong performance improvements over CNNsin tasks such as image classiﬁcation, semantic segmentation, instance segmentation, etc.
In this work we focus on image classiﬁcation using Transformers.
Fora model to do well on image classiﬁcation, it shouldlearn to only focus on the relevant information and ignore other details (eg. background information).
Selfattention does not inherently have this inductive biasof ignoring irrelevant information since it models all pairwise interactions between the inputs.
Weposit that adding a limited bandwidth temporal latent bottleneck into the Transformer will allow themodel to focus only on the most important information in the image which should enable the modelto perform well.
We test our hypothesis on the CIFAR10 and CIFAR100 [44] image classiﬁcation datasets.
We alsotest the generalization abilities of the models by comparing their performance on images of higherresolution ( 128 × 128 ) than seen during training ( 64 × 64 ).
We use ViT [21] and Swin Transformer V2(denoted as Swin V2)
[50] as our baselines.
Swin Transformer V2 has a key strength of generalizingto higher resolution images than those seen during training, making it a strong baseline.
The inputimage is split into patches of size 4 × 4 and fed in rastor order to all the models.
For the proposedmodel we use ViT as the perceptual module and add a temporal latent bottleneck module to it.
Wecall this model V I T + TLB.
To predict the classiﬁcation scores, we take the mean across the ﬁnaltemporal latent bottleneck state vectors and pass the resulting representation through an MLP.
We present the results for this experiment in table 1.
We can see that V I T + TLB outperforms V I T forall cases and performs competitively to Swin Transformer V2 .
For further hyperparameter details,we refer the reader to Appendix section 8.1.
One essential component of our model is top-down conditioning.
Top down information helps integrate information from the past into the the perceptual module.
We hypothesize that top-downconditioning enables the perceptual module to pay attention to the most important information inthe input.
We conjecture that without this, the perceptual module would have no prior knowledge ofwhich input patches to pay more attention to, hence degrading the performance of the model.
To testthis, we run an experiment where we omit top-down conditioning.
We show the results for this ablation in Table 2.
We can see that the performance drops signiﬁcantly without top-down conditioningthus showing the importance of top-down conditioning.
[Figure]
Self Supervised Learning.
Many recent works haveused vision Transformers for self-supervised learning [9, 2, 30, 12, 48, 47].
Here we show a proofof-concept that introducing a temporal latent bottleneck in Vision Transformers results in better selfsupervised representations.
We consider the SiTmodel from [2] for this experiment.
They use 3 objectives to pretrain their model - (1) The Reconstruction Objective - Reconstructs the input image, (2) The Rotation Prediction Objective - Predicts the rotation angle from [ 0 ◦ , 90 ◦ , 180 ◦ , 270 ◦ ], and (3)
The
Constrastive Objective (similar to SimCLR [14]).
For the proposed approach, we introduce a temporal latent bottleneck into SiT resulting in the SiT + TLB model.
SiT also uses additional trainable contrastive and rotation tokens as input for calculatingthe contrastive and rotation objectives respectively.
For SiT + TLB, we take the mean across the temporal latent bottleneck state vectors and use the resulting representation for computing the rotation andcontrastive objectives.
We use a chunk length of 20 for the SiT + TLB model.
We use the CIFAR10dataset [44] for our experiments.
We pretrain the model for 400 epochs and evaluate the pretrainedmodel at different epochs using linear probing.
We use the CIFAR10 dataset for linear probing aswell.
We present the results for this experiment in Figure 2.
We can see that the proposed approach outperforms SiT in all cases thus showing the effectiveness of the proposed architecture forself-supervised learning.
For additional details regarding the setup, we refer the reader to Appendixsection 8.2.
4.2 Temporal Latent Bottleneck for Sequential Decision Making
Transformers have recently been used for sequential decision making in reinforcement learning taskssuch as Atari and BabyAI [13, 37].
These works deploy Transformers in the ofﬂine RL setting wherea large number of trajectories are available either through another trained agent or an expert agent.
The Transformer is trained as an autoregressive generative model that predicts actions conditionedon the past context.
We incorporate the temporal latent bottleneck module into the Transformer andexplore its beneﬁts in the RL setting.
We test the proposed model in the BabyAI [15] and Atari [10]benchmarks.
We describe our setups in detail below.
Instruction Based Decision Making: BabyAI.
BabyAI [15] provides a suite of environmentswhere the agent has to carry out a given instruction in a partially-observable maze.
These instructions include competencies such as going to an object in the maze, placing an object beside anotherobject in the maze, opening a door with a key, etc.
Some environments in the benchmark containinstructions that combine multiple competencies sequentially.
For example, pick up a red ball andopen the door in front of you after you pick up the grey ball on your left and pick up a red box . Each environment in Baby AI benchmark has a different type of instruction that tests a differentcompetency.
The BossLevel is the most complicated environment that contains instructions from allcompetencies.
For more details regarding the various environments from the BabyAI benchmark,we refer the reader to Appendix section 8.4.
We train our models with behavior cloning using expert trajectories from an oracle.
For evaluation,we test the model by directly deploying it in the environment.
We report the success rate which mea- [Figure] sures whether the agent successfully carried out the given instruction or not.
We use a Transformer[61] as the baseline in these experiments.
For the proposed model, we introduce a temporal latentbottleneck into the Transformer-based perceptual module.
For the baseline Transformer model, weappend the language instruction to the sequence of states allowing the model to attend to the language instruction at each layer.
For the proposed model, the language instruction is appended toeach chunk, allowing each chunk to attend to it.
We consider two settings - Single task and Multi task .
In the single task setting, we evaluate theproposed approach on individual environments from the BabyAI benchmark while in the multi-tasksetting we train a single model on 8 different environments.
Single Task.
We present the results for BossLevel in Figure 3 (left) and present the results forthe other tasks in Appendix Figure 4.
We can see that while Transformer and Transformer + TLBachieve almost similar performance at convergence.
However, Transformer + TLB is much moresample efﬁcient, converging much faster .
We posit that the temporal latent bottleneck module prohibits the model from paying attention to unnecessary information which allows it to converge faster.
Multi Task.
We present the results for the multi task setting in Figure 3 (right).
We train the modelon 8 environment - PutNext, Unlock, Synth, GoToSeq, SynthLoc, GoToImpUnlock, BossLevel.
We evaluate the model on the same 8 environments.
We report the average success rate across 8games.
We can see that the Transformer + TLB model converges faster and also outperforms the
Transformer .
We refer the reader to the appendix for more details regarding the model and training. [Table]
Atari.
[13] recently introduced the Decision Transformer (DT) which learns to play variousgames in the Atari benchmark from suboptimaltrajectories of a learned agent.
Decision Transformer models the ofﬂine RL problem as a conditional sequence modelling task.
The modeluses a causal mask and supervised training tomatch the actions in the ofﬂine dataset conditioned on the future expected returns and thepast history.
This is done by feeding into themodel the states, actions, and the return-to-goˆ R c = (cid:80) C c (cid:48)
=
c r c , where c denotes the timesteps.
This results in the following trajectory representation: τ = (cid:0)
ˆ R 1 , s 1 , a 1 , ˆ R 2 , s 2 , a 2 , ˆ R 3 , s 3 , a 3 , . . .
(cid:1) , where a c denotes the actions and s c denotes the states
At test time, the start state s and desired return ˆ R 1 is fed into the model and it autoregres- sively generates the rest of the trajectory
Experimental results show that DT can leverage the strong generalization capabilities of Transformers and achieve the desired returns in a wide variety of tasks in Atari and OpenAI Gym, outperforming previous approaches in ofﬂine RL
We use the same setup as used in [13] for our experiments.
We set the context length to a ﬁxednumber C .
During training, C timesteps from an episode are sampled and fed into the modelresulting in a trajectory of length 3 C (considering 3 modalities - returns-to-go, states, and actions).
Each modality is processed into an embedding of size d .
The state is processed using a convolutionalencoder into an embedding of size d .
The resulting trajectory is fed into the decision Transformer.
The outputs corresponding to the states s c are fed into a linear layer to predict the action a c to betaken at timestep c .
For the proposed model, we incorporate a temporal latent bottleneck moduleinto the decision Transformer.
We present our results in Table 3.
The model is trained on 1% of the Atari DQN-replay dataset [1](500K transitions for each game).
We use the same 4 games used in [13]: Pong, Seaquest, Qbert,and Breakout.
We can see that the proposed model outperforms decision Transformer in all theconsidered games thus showing the effectiveness of the proposed model .
More details regarding themodel and training can be found in the appendix section 8.5.
4.3 Temporal Latent Bottleneck for Language Modeling
We also apply the Transformer + Temporal Latent Bottleneck Model to language modelling on theenwiki8 dataset [52].
We use Transformer XL [19] as our baseline.
We ﬁnd that Transformer XLachieves a performance of 0.99 bit-per-byte while the proposed model achieves 0.97 bit-per-byte(lower is better).
Thus showing that the proposed model is useful in language modeling as well.
sults averaged across 3 seeds. [Table] Temporal Latent Bottleneck For Long Range Dependencies.
Here, we test the effectiveness of the proposed model in modellinglong range dependencies.
We apply the proposed model on the ListOps and text classiﬁcation tasks from the Long Range Arena (LRA)benchmark [60].
Both these tasks have verylong sequences ranging from 1K to 4K tokens.
[Table] Thus, for a model to do well, it has to learnto capture dependencies across very long timescales.
Additionally, all these tasks have aninherent hierarchical structure.
For example, Listops consists of nested list operations whichmakes it hierarchical.
For text classiﬁcation,the inputs consist of text in the form of bytes.
Therefore, the model has to learn to compose bytesinto characters and characters into words.
We hypothesize that the multi-scale nature of the proposed model will be extremely useful in modelling such hierarchical information.
The temporallatent bottleneck which operates at a slow scale can behave as a composer that composes low-levelinformation in relevant way to solve the desired task.
For this experiment, we use the same setup as [68].
For the proposed model, we use a Transformer asthe perceptual model and implement the temporal latent bottleneck as described in Section 2.2.
Wetake the mean across the temporal latent bottleneck state vectors and use the resulting representationfor classiﬁcation.
We compare the model against the long-short Transformer (LS) model [68], whichis a recently proposed model for the long range arena benchmark, and the vanilla Transformer model[61].
We present the results in Table 4.
We can see that the proposed model outperforms both thebaselines in both the tasks thus showing its usefulness in modeling long range dependencies .
Forfurther details, we refer the reader to Appendix section 8.3.
Conclusion
We have developed an approach aimed at introducing selectivity in the interactions across time-stepsin a transformer by splitting processing into two streams: (a) a slow stream that is updated in a recurrent manner and (b) a fast stream that processes the visual input.
The two streams are parameterizedindependently and interact with each other via attentional bottleneck.
The information processed bythe fast stream is used to change the state of the slow stream, and the information in the slow streamis used by the fast stream as contextual information to process the input.
Through our experimentswe show that the proposed approach works well across wide range of domains and problems.
Onelimitation of the proposed model is that the chunk size is ﬁxed and treated as a hyperparameter whichrequires some domain knowledge.
Future work should explore methods for dynamic chunking Checklist
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes](c)
Did you discuss any potential negative societal impacts of your work?
[Yes](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] Appendix Section
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes](b) Did you mention the license of the assets?
[No] All the datasets we use are openlyavailable
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data
you’re using/curating?
[N/A]
The datasets we use do not contain any sensitive in-
formation.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁ-
able information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
[Table]
a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable
N/
[Table]
you describe any potential participant risks, with links to Institutional Review
[Table]
approvals,
[Table]
you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation
N
A]
[Table]