保留正文与补采附件。源页面当前未重新核验;具体限制见交付清单。
We value your privacy
We use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking "Accept All", you consent to our use of cookies.
Customise Reject All Accept All
Customise Consent Preferences
We use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below.
The cookies that are categorised as "Necessary" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more
NecessaryAlways Active
Necessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.
Cookie BCTempID
Duration 10 minutes
Description No description available.
Cookie __cf_bm
Duration 1 hour
Description This cookie, set by Cloudflare, is used to support Cloudflare Bot Management.
Cookie AWSALBCORS
Duration 7 days
Description Amazon Web Services set this cookie for load balancing.
Cookie _cfuvid
Duration session
Description Cloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services
Cookie li_gc
Duration 6 months
Description Linkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes.
Cookie __hssrc
Duration session
Description This cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session.
Cookie __hssc
Duration 1 hour
Description HubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie.
Cookie wpEmojiSettingsSupports
Duration session
Description WordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly.
Cookie BCSessionID
Duration 1 year 1 month 4 days
Description Blueconic sets this cookie as a unique identifier for the BlueConic profile.
Cookie _octo
Duration 1 year
Description No description available.
Cookie logged_in
Duration 1 year
Description No description available.
Cookie __Secure-YEC
Duration past
Description YouTube sets this cookie to stores the user's video player preferences using embedded YouTube video
Cookie __eoi
Duration 6 months
Description Description is currently not available.
Cookie AWSALBTGCORS
Duration 7 days
Description No description available.
Cookie login-status-p
Duration past
Description Description is currently not available.
Cookie AWSALBTG
Duration 7 days
Description No description available.
Cookie csrf_token
Duration session
Description No description available.
Cookie token_v2
Duration 1 day
Description Description is currently not available.
Cookie D
Duration 1 year
Description Description is currently not available.
Cookie PHPSESSID
Duration session
Description This cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed.
Cookie VISITOR_PRIVACY_METADATA
Duration 6 months
Description YouTube sets this cookie to store the user's cookie consent state for the current domain.
Cookie cookietest
Duration session
Description The cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience.
Cookie __Host-airtable-session
Duration 1 year
Description This cookie is used to enable us to integrate the services of Airtable.
Cookie __Host-airtable-session.sig
Duration 1 year
Description This cookie is used to enable us to integrate the services of Airtable.
Cookie m
Duration 1 year 1 month 4 days
Description Stripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly.
Cookie BIGipServer*
Duration session
Description Marketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user.
Cookie __cfruid
Duration session
Description Cloudflare sets this cookie to identify trusted web traffic.
Cookie _GRECAPTCHA
Duration 6 months
Description Google Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks.
Cookie __Secure-YNID
Duration 6 months
Description Google cookie used to protect user security and prevent fraud, especially during the login process.
Cookie cookieyes-consent
Duration 1 year
Description CookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors.
Functional Functional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.
Cookie lidc
Duration 1 day
Description LinkedIn sets the lidc cookie to facilitate data center selection.
Cookie brw
Duration 1 year
Description No description available.
Cookie brwConsent
Duration 5 minutes
Description Description is currently not available.
Cookie WMF-Uniq
Duration 1 year
Description Description is currently not available.
Cookie loom_anon_comment
Duration 1 year
Description No description available.
Cookie loom_referral_video
Duration session
Description Description is currently not available.
Cookie VISITOR_INFO1_LIVE
Duration 6 months
Description A cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface.
Cookie yt-remote-connected-devices
Duration Never Expires
Description YouTube sets this cookie to store the user's video preferences using embedded YouTube videos.
Cookie ytidb::LAST_RESULT_ENTRY_KEY
Duration Never Expires
Description The cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future.
Cookie yt-remote-device-id
Duration Never Expires
Description YouTube sets this cookie to store the user's video preferences using embedded YouTube videos.
Cookie yt-remote-session-name
Duration session
Description The yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video.
Cookie yt-remote-fast-check-period
Duration session
Description The yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos.
Cookie yt-remote-session-app
Duration session
Description The yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player.
Cookie yt-remote-cast-available
Duration session
Description The yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player.
Cookie yt-remote-cast-installed
Duration session
Description The yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video.
Cookie cp_session
Duration 3 months
Description Codepen sets this cookie for Help systems found in the website.
Cookie loid
Duration 1 year 1 month 4 days
Description This cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform.
Analytics Analytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.
Cookie __hstc
Duration 6 months
Description Hubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session).
Cookie hubspotutk
Duration 6 months
Description HubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts.
Cookie _ga
Duration 1 year 1 month 4 days
Description Google Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors.
Cookie ga*
Duration 1 year 1 month 4 days
Description Google Analytics sets this cookie to store and count page views.
Cookie __Host-psifi.analyticsTrace
Duration 6 hours
Description Description is currently not available.
Cookie __Host-psifi.analyticsTraceV2
Duration 6 hours
Description Description is currently not available.
Cookie _gh_sess
Duration session
Description GitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form.
Cookie YSC
Duration session
Description YSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages.
Cookie ajs_anonymous_id
Duration 1 year
Description This cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before.
Cookie vuid
Duration 1 year 1 month 4 days
Description Vimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website.
Performance Performance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.
Cookie AWSALB
Duration 7 days
Description AWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target.
Cookie acq
Duration past
Description Description is currently not available.
Cookie acq.sig
Duration past
Description Description is currently not available.
Cookie ptc
Duration 2 years
Description No description available.
Advertisement Advertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.
Cookie muc_ads
Duration 1 year 1 month 4 days
Description Twitter sets this cookie to collect user behaviour and interaction data to optimize the website.
Cookie guest_id_marketing
Duration 1 year 1 month 4 days
Description Twitter sets this cookie to identify and track the website visitor.
Cookie guest_id_ads
Duration 1 year 1 month 4 days
Description Twitter sets this cookie to identify and track the website visitor.
Cookie personalization_id
Duration 1 year 1 month 4 days
Description Twitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting.
Cookie guest_id
Duration 1 year 1 month 4 days
Description Twitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences.
Cookie bcookie
Duration 1 year
Description LinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs.
Cookie __Secure-ROLLOUT_TOKEN
Duration 6 months
Description YouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment.
Cookie yt.innertube::nextId
Duration Never Expires
Description YouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.
Cookie yt.innertube::requests
Duration Never Expires
Description YouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.
Cookie session_tracker
Duration session
Description This cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website.
Cookie edgebucket
Duration session
Description Reddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content.
Cookie did
Duration 1 year
Description Arbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year.
Uncategorised Other uncategorised cookies are those that are being analysed and have not been classified into a category as yet. No cookies to display. Reject All Save My Preferences Accept All Toggle SearchToggle Mobile Navigation Publish AI, ML & data-science insights to a global community of data professionals. Sign inSubmit an Article LatestEditor's PicksNewsletterShipAI
Write For TDS LinkedInXToggle Search Deep Learning
Part 3 - Towards Mamba State Space Models for Images, Videos and Time Series Sascha Kirch September 3, 202418 min read
Image by Sascha Kirch.
This is part 3 of my new multi-part series 🐍 Towards Mamba State Space Models for Images, Videos and Time Series. Mamba, the model to be said to replace the mighty Transformer, has come a long way from the initial idea of using state space models (SSMs) in deep learning. Mamba adds selectivity to state space models which results in Transformer-like performance while maintaining the sub-quadratic work complexity of SSMs. Their efficient selective scan is 40x faster than a standard implementation and can achieve 5x more throughput as a Transformer. Join me on this deep dive on Mamba where we will discover how selectivity solves the limitations of previous SSMs, how Mamba overcomes new obstacles that come with those changes and how we can incorporate Mamba into a modern deep learning architecture. Mamba State Space Models for Images, Videos and Time Series
Learn this step by step with the interactive Machine Learning roadmap.
Structured State Space Models Visually Explained We ended Part 2 with the introduction of the S4 model. By constraining the state matrix A to have a certain structure, the convolution kernel could be computed very efficiently making S4 30x faster using 400x less memory compared to previous SSMs, hence making it a feasible alternative to Transformers. We said that our structured state space models can be efficiently trained and efficiently inferred due to the duality of the recurrent and convolutional representations of LTI-Systems. ❗ The issue though: An LTI-System treats all its inputs equally and therefore cannot make data dependent decisions.
Mamba addresses this problem by introducing a data dependent selection mechanism, but as we will see this comes at a cost. Later more about that. Let's start with some intuition.
A fundamental problem of sequence modelling is to decide to which extent the context (meaning previous inputs) is compressed into the state of the model.
Remember that the RNN achieves fast inference due to its highly compressed state that models long-range dependencies ineffectively. A Transformer on the other hand, explicitly does not compress the context at all. Each token can attend to each other token at any time, allowing very effective modelling of long-range dependencies at the cost of quadratic work complexity O(L²).
The Structured SSM approximates the context using the HiPPO matrix resulting in some compression, while it can be trained more efficiently as the RNN because of its convolutional representation.
Figure 1: Efficiency vs. performance of the RNN, SSM, Transformer and the Mamba model. Image by Sascha Kirch.
What we want is something in between, or let's better say the best of both worlds. A performant and efficient model. And spoiler alert: Mamba might be a promising candidate.
Instead of attending to the entire history like the Transformer and instead of compressing the entire history into a single state like RNNs and SSMs, we want to selectively compress just the relevant data to end up with a highly performant and efficient model. Efficient in this context means not only inference speed but also relates to other resources like memory and data needed to train the model.
Let's look at this simply example of predicting my name:
Figure 2: Selectively put focus on a subset of inputs. Image by Sascha Kirch.
Up until now all matrices A , B and C have been independent of the input data. They did not change for different time steps of the signal, hence we said we have a linear and time-invariant (LTI) system where the recurrent and convolutional representation were equal.
Figure 3: Different representations of state space models. Image by Sascha Kirch.
Mamba changes this by making matrices B and C as well as the discretization parameter ∆ time dependent. This means for every input token we end up with a different set of parameters! Recall that matrix A is a structured matrix derived from the HiPPO matrix and approximates how the state evolves over time based on the history of previous states. Mamba keeps it that way.
Making the SSM selective introduces new dimension into the matrices. Explicitly, we now need to consider each token from a sequence of length L parallelized to a batch size B.
Figure 4: Tensor shapes of all matrices involved in the selective state space model. Image by Sascha Kirch.
There are a few things to note here:
❓ But, what does all this even mean? What is the intuition behind all those matrices and why are only some selective?
Matrix A: The state-to-state matrix decides what from the previous state propagates into the next state. Matrix B: the input-to-state matrix transforms any incoming token. It must be selective to filter out irrelevant data or to focus on important data to allow only relevant data to enter into the new state. It resembles some sort of gating mechanism. Matrix C: The state-to-output matrix transforms the state to an output token. It must be selective to decide which information from the state is required to materialize the output. Hence, like matrix B , matrix C has also some connections to a gating mechanism. Matrix Δ: The discretization parameter is often referred to as the step-size of converting the continuous input signal into a discrete sequence. While in a classical SSM, this might be intuitive, in a deep learning setup it is not, especially when working with already discretized data or other modalities like text tokens. So, let's look at the extreme cases: if Δ→∞ the system focuses more (or longer) on the input, hence in the state equation the state update is dominated by the contribution of the input and less by the previous state. If Δ→0 , the opposite happens. That means that Δ can be seen as some sort of weighting between the importance of new samples compared to the previous state.
Get an email whenever Sascha Kirch publishes 🚀 _Get an email whenever Sascha Kirch publishes 🚀 Looking to learn more about deep learning or simply stay up to dat_e...medium.com
Again, we must talk about the state matrix A. Have you noticed something when looking at the dimension of matrix A? We were always talking about matrix A being a square matrix of NxN , but in Mamba its of shape DxN. What is happening?
Well, a lot has happened between the S4 paper and the release of Mamba. In particular, diagonal state space (DSS) models further simplified the DPLR (diagonal plus low-rank) formulation of S4 to be a simple diagonal matrix, by removing the low-rank term. When we talked about S4 in part 2 of this series, we said having a diagonal matrix A for the calculation of the convolution kernel K would actually be the ideal case, since it reduces the number of computations drastically!
Figure 5: Truly diagonal state matrix A as proposed by the DSS paper. Image by Sascha Kirch.
In a follow-up paper of S4's authors, they provided a thorough analysis of the effects on why this works and further suggests multiple new initialization schemes for the true diagonal case resulting in the S4D model.
Long story short, Mamba uses the S4D-Real initialization, mentioned in a side note of the S4D paper, because it has the same spectrum as the original S4-LegS HiPPO Matrix. We will further talk about initialization in chapter 5.
But now let's come back to Mamba's matrix A. Remember from part 2 of the series that if we have a diagonal matrix of NxN , we only need to store N values, meaning we can store it as a vector. The DxN matrix of Mamba is really just D times such a vector representing the diagonal matrix!
Figure 6: Creating the diagonal matrices of the SSM from Matrix A. Image by Sascha Kirch.
This allows us to observe another interesting fact about Mamba and previous state space models: Each channel in D is processed independently which can be interpreted as each channel gets its own SSM! This fact becomes important in part 6 and part 7 when we look into the MambaMixer and Mamba-2 respectively.
Having different parameters for different timesteps implies that our model is now time-variant. Hence the convolution representation is not valid anymore and we loose the benefit of fast training. A big portion of the Mamba paper is about how to overcome this problem. They introduce three strategies resulting in their selective scan algorithm : (1) parallel associative scan, (2) kernel fusion and (3) recomputation of gradients.
Note: These ideas are not completely new. The S5 layer (Simplified State Space Layer for Sequence Modelling) applied the parallel scan and H3 (Hungry Hungry Hippos) already worked with kernel fusion to speed up the training.
Recall that the structured state space model is a time invariant system (LTI) since the matrices involved in updating the state (A and B) are identical for each new input. That allowed either using a recurrent representation or a convolution representation. In Mamba this changes since the discretized matrices A and B are now dependent on the input.
That implies that we cannot use the convolution representation to speed up the training and we are left with the recurrent representation, this time with different matrices for each new input and update of the state:
Figure 7: Recurrent representation of a selective state space model. Image by Sascha Kirch.
The recurrent execution has a work complexity of O(L) , meaning doubling the length L of the input sequence requires twice as much work to compute the final output. The one big issue with recurrency is that during training we don't want to calculate all the states sequentially, because we have access to the labels of all inputs beforehand. During inference, especially for an autoregressive sampling task (e.g. next word prediction) this would be totally fine though.
❓ The question is, how can we speed up the sequential update of the state, where each state strongly depends on the previous one? It's not like that we could run it in parallel, right?
This particular type of recurrency, where each new state is the sum of the previous state and an input, is called scan. It is such a common problem in science and engineering, that even before we had parallel computers, in 1986 and 1990, techniques to parallelize a scan have been already been published. It wasn't until 2005 that these algorithms would be tested on a real GPU.
Figure 8: Parallel scan algorithm. Image by Sascha Kirch.
The parallel scan has a work complexity of O(log L) , which is smaller than O(L). The longer the input sequence L , the higher the speed-up compared to the recurrent execution.
Recall that a vanilla transformer has a work complexity of O(L²).
As you might know, any computer has different types of memory. Generally speaking, the farther away from the processing unit, the larger the memory and the more time it takes to read the data (mainly due to properties of the PCB and other electronic components that do not allow higher frequencies).
What usually happens in a GPU is that variables are read from the larger but farther away HBM (high bandwidth memory) into the faster SRAM to perform a calculation and then the result is written back into the HBM. If a calculation requires multiple steps, this process of reading, processing and writing is repeated until the entire calculation is done. If a calculation requires less time for the actual calculation as for reading and writing, it is called memory bounded aka. I/O bounded.
Figure 9: GPU utilization vs. memory I/O. Image by Sascha Kirch.
I/O bound operations use their computing hardware quite inefficiently since they spend most of the time waiting for data in an idle state. Basically, you want to avoid reading large amounts of data too frequently.
❓ But what can we do? Or better said: what did the authors of Mamba do?
Instead of storing the discretized matrices A and B , only the continuous-time matrices are stored in global memory. Matrices A , B , C and Δ are read from the HBM into the SRAM and then the discretization is performed in SRAM.
Here an example on how many values need to be read when either reading the continuous-time matrices or when reading the discretized matrices considering the shapes of the matrices as depicted earlier:
Figure 10: Number of values needed to be read for discretized matrices vs. continuous matrices. Image by Sascha Kirch.
Further, kernel fusion is applied to reduce the number of reads and writes from and to the HBM. In particular, they fuse the discretization, the selective scan and the multiplication with matrix C into a single kernel. Hence, they read all the data needed for an entire forward pass through all states, process those data and write back the final output y of shape BxLxD. The large states of shape BxLxDxN are not materialized to avoid memory requirements and reading large amount of data.
Figure 11: Impact of kernel fusion on the GPU utilization. Image by Sascha Kirch.
Having heard about kernel fusion and the different memories involved allows us to finally have a look on the iconic illustration from the Mamba paper showing the interaction of the different matrices to calculate the new states and the outputs:
Figure 12: Selective State Space Model with Hardware-aware State Expansion. Image by A. Gu et.al., annotated by Sascha Kirch.
Green indicates that those values are stored in the HBM while orange indicates that they are kept in the SRAM. Later in chapter 5 we will adapt this image to show what really happens and to find the trained parameters.
Previously we saw that with kernel fusion we speed up inferring the model through a reduction of I/O operations and the fact that the hidden states are not saved in the HBM. However, to calculate the gradients and to back propagate the loss for the weight update, one needs those intermediate states. Since the kernel fusion is so effective, recomputing the intermediate states on demand as they are needed during back propagation is much more efficient than writing and reading the large intermediate states of shape BxLxDxN from the HBM.
Having spent so much effort on making SSMs faster, one question remains: how fast is it compared to the most efficient attention implementation we know: FlashAttention-2.
Figure 13: Efficiency Benchmarks of Mamba. Image by A. Gu et.al., annotated by Sascha Kirch.
We are now equipped with the selective SSM, a state space model that is efficient, fast, scales linearly with the input sequence length L and that is also selective since each input is processed with its own set of parameters. The only thing left for us is to do is to wrap it into a block and stack multiple blocks on top of each other to build a deep learning model; the same way as we would stack the multi-headed Transformer block in a Transformer-Decoder.
Let's see what a model for an autoregressive next word prediction using Mamba could look like:
Note that the names of dimensions and names for the modules are the same as in the official code implementation of the Mamba Block to better understand how it can be configured and what is happening.
Figure 14: Detailed block diagram of the Mamba-Block. In this example: L=4, d_model=3, expand=2, d_inner=6. Image by Sascha Kirch. The input projection simply projects each of the L tokens into a
d_model-dimensional space. The head does the prediction using one or more output tokens from the stack of Mamba Blocks. Inside the Mamba Block , we first add theresidualsfrom the previous layer to thehidden_statesand then output their sum as the nextresiduals. We normalize using RMS Layer Normalization and then linearly project fromd_modeldimensions into2*d_innerdimensions, whered_inneris defined asexpand*d_model. Inside the codeexpanddefaults to 2. This output ofin_projis then split into 2 different paths, both of which are now of dimension[1, d_inner]and both containing a SiLU non-linearity. The Selective SSM block is our state space model with the known state equations parametrized by the matrices A , B , C and Δ. Note that the vector dimension D from Figure 4 corresponds tod_inner. After the Selective SSM block, the two paths are combined, projected back fromd_innertod_modelto form the final output. Note that both linear projections are broadcasted over the L tokens, meaning they use the same parameters for each of the tokens!
While inside the Mamba Block it is pretty clear where the learned parameters are located (i.e. input projection in_proj, output projection out_proj and the 1D convolution conv1d) it is still a mystery where the learned parameters hide inside the selective SSM Block.
I think we can best observe what's going on if we adapt Figure 12, to show what really happens inside the selective SSM (single token and no batch axis):
Figure 15: Selective State Space Model with Hardware-aware State Expansion including the discretized matrices and the learnable linear projections. Image by Sascha Kirch.
Again, a lot is going on this figure, so let's start with two observations:
Now let's have a look on what is happening. Note that the time-dependent matrices B and C are the result of the linear projection x_proj of the input x at time t. Again, a single linear layer is used, likely for performance reasons. The output of x_proj is further projected from dt_rank to d_inner to form the final discretization matrix Δ. It is only Matrix A we optimize directly (In reality we keep track of log(A) for numerical stability during training).
To conclude, inside the Selective SSM, we do not train the matrices used in the state equations, but we train the weights of linear projections that result in those matrices. It is only Matrix A we are optimizing directly.
And with that, there is only one question left to answer...
Not an easy question to be answered, since not only reading the papers is complicated due to the complicated matter of the subject, but also the code is quite convoluted.
The first thing to notice: There is the initialization of a stand-alone Mamba Block and then there is a reinitialization, depending on the final architecture of the deep learning model that stacks multiple blocks together i.e. by using the [_init_weights()](https://github.com/state-spaces/mamba/blob/62db608da60f6fc790b8ed9f4b3225e95ca15fde/mamba_ssm/models/mixer_seq_simple.py#L85-L115) function.
Let's fist look at the Mamba Block itself.
All modules are initialized with pytorch's default initializer of the respective type of module (e.g. uniform for Conv1D and Linear) with two exceptions: Matrix A and dt_proj, the linear projection used to construct Δ.
Matrix A is relatively straight forward to initialize using the S4D-Real initialization mentioned earlier:
Figure 16: S4D-Real initialization of state matrix A. Image by Sascha Kirch.
The timescale on the other hand is a bit more complicated. I recommend you to check out the "How to Train your Hippo" paper for more details. The short version is:
The bias of dt_proj shall be initialized by log-uniformly sampling in a range of (Δmin=0.001, Δmax= 0.1). A very important side note in that paper that is easy to oversee: You can increase the performance if you set 1/Δmin to be approximately in the range of the sequence length. So 0.001 would be approximately 10³ samples.
For example, look at this plot from the paper on a benchmark dataset with a sequence length of L=16,384 :
Figure 17: Impact of choice of timescales on validation accuracy on Path-X benchmark. Image by A. Gu et.al., annotated by Sascha Kirch.
Finally, and without going too much into detail, here is the[_init_weights()](https://github.com/state-spaces/mamba/blob/62db608da60f6fc790b8ed9f4b3225e95ca15fde/mamba_ssm/models/mixer_seq_simple.py#L85-L115) function with some annotations that is used to overwrite already initialized weights when they are used inside the MixerModel and the MambaLMHeadModel the authors provided:
This function is actually called in both those classes meaning they overwrite the overwritten weights from the initial initialization 😵 But jokes aside, if working under time pressure and finding the most performant solution there is not always enough time to write clean code. Yet another trade-off...
Figure 18: Annotating the _init_weights() function from the Mamba repository that is used to overwrite model weights after they have been initialized. Code taken from here, annotated by Sascha Kirch.
With Mamba we now end up with a sequence model as powerful as the Transformer and as efficient as an RNN. We saw that while we lost the duality of the recurrent and the convolutional representation by adding selectivity, we were able speed up the model again by implementing the selective scan algorithm. The transformer is a general-purpose architecture, applied to all data modalities you can think of. So far, when talking about Mamba, we were only talking about sequence-like data such as text tokens or analog signals. ❓ But, can we apply Mamba to non-sequence data? Let's say images?
🤝 I'd like to thank you, t he readers of my stories, for your constant support, clapping for my stories and, following me to never miss any of my latest articles. It is you who motivates me to continue writing, dive deep into complex topics, and become a better engineer/researcher by writing about them. You are awesome!
Vision Mamba: Like a Vision Transformer but Better In part 4 we will see how Mamba-like models can be applied on non-sequential data types such as images. We'll revisit certain design choices of the Vision Transformer (e.g. patch size, image resolution and position encodings) and see if they are still valid when using Mamba instead of attention. Vision Mamba ends up being 2.8x faster than DeiT and saves 86,8% GPU Memory, and we will see how this is possible. So stay tuned!
Mamba State Space Models for Images, Videos and Time Series Paper Walkthroughs by Sascha Kirch
All articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more. Register for freeSign in to your account
Written By
See all from Sascha Kirch Deep DivesDeep LearningMamba Image Video SignalParallel ComputingState Space Model Share This Article
Towards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program. Write for TDS
Artificial Intelligence An experiment in pure AI development that challenges everything you've heard about coding assistants. Daniel Bentes February 19, 202510 min read
Artificial Intelligence How to use BERT from the Hugging Face transformer library for four important tasks Saketh Kotamraju January 18, 202112 min read
Data Science Solve a puzzle for fun. Learn a new Algorithm. Win-win. Full code. Naresh Ram April 10, 202211 min read
Artificial Intelligence Triton is a language compiler to write highly optimized CUDA kernels. Learn the basics of GPU programming and how to use Triton. Kushajveer Singh September 12, 20218 min read
Deep Learning The theory of Sound and Acoustics could be used for cutting edge Emotion-Recognition technologies. here's how we can apply it using Python Tal Baram February 25, 202110 min read
Science An In-Depth Interview with Professor Coello Coello Emil Rijcken October 6, 202313 min read
Artificial Intelligence Learn how to compare classifiers using hypothesis testing Marcello Politi July 6, 20225 min read
Deep Learning Steps you can use to pass, too - even if you got an M1 Mac Pascal Janetzky July 31, 20217 min read
Artificial Intelligence How does it help? NVS Yashwanth September 18, 20204 min read
Your home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals. Subscribe to Our NewsletterWrite For TDS•About•Advertise•Privacy Policy•Terms of Use © Insight Media Group, LLC 2026•The New Stack•roadmap.sh•Creator Network word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word word
mmMwWLliI0fiflO&1
mmMwWLliI0fiflO&1
mmMwWLliI0fiflO&1
mmMwWLliI0fiflO&1
mmMwWLliI0fiflO&1
mmMwWLliI0fiflO&1
mmMwWLliI0fiflO&1
Close