How Moonshot Makes Open Models Great

Hi, everyone.

Thank you so much for the introduction.

It's great to be here.

It's great to have this opportunity to share with you guys some of our latest progress and efforts.

So one of our major pursuits is to build better open models.

And we believe in democratizing intelligence.

With open models, you can deploy anywhere.

You can be on your local servers.

You can be on the cloud.

And you can access every single bit of the weights in the modeling instead of just using a black box.

And this is one of the slides that I took from Jensen's talk earlier this year at CES.

So as you can see, open models are quickly closing the gap with proprietary models.

And it's reaching the frontier.

And we believe that with better and better open models, we're going to make intelligence more accessible to anybody in the world in every corner of the world.

But open models cannot be just open.

They have also to be great.

So in this talk, we're going to discuss how we make open models great.

So as we know, scaling is a primary driver for a lot of progress, maybe all of the major AI developments that we have witnessed in the last few years.

And here, we're going to discuss how we scale our model in different dimensions.

So on the left-hand side, the first figure you see here is kind of the standard scaling law.

So on the x-axis, you have the log of the number of training tokens.

And on the y-axis, you have the log loss.

And as you scale the number of training tokens, you get a lower loss.

But here, the point is, we're not going to just scale the number of training tokens.

But we also want to improve the token efficiency, meaning that we want to move this curve to the left-hand side so that we can achieve a lower loss, a much lower loss, using the same number of training tokens.

And this can be achieved by having better architectures and optimizers, as we'll discuss in our later slides.

And the second scaling dimensions that we're very interested in is to scale the context length.

So as you can see in the second figure, if we increase the context length, then we can have a much higher accuracy in terms of predicting the token loss at a given position.

And this means that we can increase the capability of the model to achieve more complex tasks by increasing the context length.

So this is the second scaling dimensions that we're going to talk about.

And the third scaling dimension is the number of agents.

So we introduced this new learning paradigm of agent swarms, where we don't just rely on a single agent, but we also orchestrate a swarm of agents that can accomplish the subtasks in parallel so that we can increase the task capacity.

And we can translate all of this into the language of agents.

So if you look at token efficiency, it's mostly about having a stronger prior so that you can have more efficiency when you do agent RL to search for a better solution.

And we think about long context, it's mostly about increasing the context length so that you can have a longer running agent.

It can probably run for days or even weeks or months to accomplish more tasks, more complex tasks.

And for agent swarms, it's another dimension that added it.

And at the end of the day, we're going to have a swarm of agents that each of them have a super long context, and each of them have a very strong prior for us to search in this entire agent RL system.

So we're going to start from token efficiency.

So this is one of the most classical figures in the history of machine learning.

So it's taken from Kaplan et al.

And it basically says that if we scale proportionally the number of training tokens, the model parameters, and also the amount of compute, we can get lower and lower loss.

And this is one of the major breakthroughs that the entire community has achieved in the last few years to get better intelligence.

But here, what we're interested in is to have better and better token efficiency.

And here's the thing.

So one thing that I would like to emphasize is that token efficiency is not just about efficiency.

It's actually also about improving the upper bound of intelligence.

So here's why.

So suppose you have, say, 50 trillion tokens, 50 trillion high-quality tokens.

And then you apply this new optimizer, maybe the MUN optimizer.

And then all of a sudden, you have a two-time token efficiency.

So it means that it's almost like magic that you get equivalently 100 trillion tokens.

And nowadays, we are scaling towards the data war.

And we're hitting the data war.

And the amount of high-quality data is quite limited.

And if we suppose that it's a constant amount, then we increase the token efficiency.

It means that we're going to get better intelligence out of it.

It's not just about infrastructural efficiency.

It's about better intelligence.

So this is why we spend a lot of efforts in this aspect, because it's going to push the frontier of intelligence.

And MUN optimizer is one of the things that we have heavily invested in since last year.

So it's a second-order optimizer.

And basically, every single gradient update is transformed in a way that each entry is orthogonal to each other.

And this is very different from the traditional Adam optimizer.

And if we implement this optimizer properly, you can get a two-times token efficiency improvement.

So we are the first work.

We published the first work to demonstrate that MUN optimizer is actually scalable for LLM training.

And these are two key techniques that we employ to make it effective for large-scale training.

So one of them is weight decay.

It is critical for scaling to larger models.

And the second is we want to ensure a consistent RMS updates compared to Adam.

So we have this adjustable coefficient that is applied to each update so that the resulting RMS is going to be comparable to Adam.

And to make MUN memory efficient across all these NVIDIA GPU clusters, we also developed a distributed MUN optimizer implementation that partitions the states across the data parallel group so that we can have a very efficient implementation for the MUN optimizer.

And these are some of the results that were presented in the paper.

So as you can see, with the same number of parameters and the same number of training tokens, we just replaced the original Adam W optimizer with the new MUN optimizer.

It's going to improve the performance across the board significantly.

But there was this new challenge that we encounter when we try to scale it up further.

When we try to scale MUN for a one trillion parameter model, we encounter a new issue about training instability.

So as you can see on the left figure, we observed that the max logits quickly explodes and quickly exceeds 1,000.

And the typical values for training for this max logits is about, say, 50 or maybe less than 100.

But for MUN, it quickly exceeds 1,000.

And at the same time, we observe training divergence on the left-hand side.

If you look at the training loss, it goes down a bit.

But then at the end of the day, it explodes.

And it cannot converge as expected.

So this is one of the technical challenges that we have to address.

And the solution to this is to introduce this new technique called QKclip.

So basically, what it says is that for each attention had in this entire neural network, in the forward pass, we're going to compute the max logits.

And then we're going to calculate a dividing factor that can be applied to each key projection, as well as the query projection, so that we can clip the maximum value of the query and the key to constrain it into a given range.

So that we're not going to have explosion anymore.

So these are some of the empirical results.

On the left-hand side, there are two curves.

But they are strictly overlapped with each other.

So these are the training curves before and after applying the clipping technique.

So you can see, the clipping technique does not affect the training loss decrease at all.

But on the right-hand side, if we inspect the intermediate metric, if we inspect the max logit, it's going to be effectively constrained.

So if it's exposed, as before, but at the value of 100, it's going to be clipped at a constant value for a long time.

And then after a certain number of steps, it will just naturally go down.

So the neural network sort of find a way to constrain the maximum value of the max logit to ensure a stable training process.

And at the same time, it doesn't affect the training convergence as shown in the left figure.

So we employed this technique in our K2 model training and successfully scaled it to 1 trillion parameters.

And this is the first example of a large-scale muon training in the history of machine learning.

And the second dimension that we're very interested in is long context.

So this is another figure.

It's probably less known.

It's one of the hidden gems in these papers.

So instead of just pushing down the training loss by training on more tokens, it has some insights from another perspective.

So as we can see, this is a comparison between transformers and LSTMs.

So on the left-hand side, we can see that transformers achieve a lower training loss given the same number of parameters and the same number of training tokens as expected.

And this is why transformers become the sort of the de facto architecture that people are using right now.

But on the right-hand side, it's very interesting to see that transformers are actually better because it can improve through the whole context.

So the x-axis is the token index in context.

And if you increase the token index, you can see that the training loss of transformers actually dropped by a lot.

If you just continually increase the context length, the loss just continuously drops down.

But if you look at the curve of LSTM, it's saturated after a certain number of tokens.

It means that transformers have this better capability of capturing longer context.

And this is what makes it better because if it go back to like 10 years ago, people used LSTM for tasks like machine translation.

But it is not good for, for example, understanding the entire code base or running a super long agent trajectories to solve, for example, writing Linux kernels from scratch.

It's not going to be accomplishable by LSTMs.

So this is a very much-needed capability in the era of agents because tasks are becoming harder and harder, and we need longer and longer context.

So the research idea here is to develop a better architecture so that we can efficiently scale to a longer context length and, at the same time, achieve a lower per token loss at larger token indices.

And this is the motivation for which we introduce this new architecture called KimiLinear, and it contains this new linear attention variant called KimiDeltaAttention, which improves the original gated delta rule, GDR, by improved recurrent memory.

I will show the details later.

And at the same time, we're going to mix linear attention layers with full attention layers using a 1-to-3 ratio so that you can balance between these long context capabilities and, at the same time, having a more efficient implementation.

So this is some of the formulation.

The basic idea is simple.

If we look at linear attention, in the original formulation, the memory is going to be global.

So there is a global single decay factor that is applied along the way.

So it means that, basically, there are only two cases.

In one case, you're going to forget basically everything, and you're not going to retain any information.

In the second case, you can choose to retain almost everything, but at the same time, you don't have the capability to leave out some of the unnecessary information in this long context.

So we introduce this key idea of having a fine-grained decay factor, as shown in this highlighted alpha term.

So instead of being a scalar, it's going to be a diagonal matrix, which controls the decay rate for each channel so that we can have two possibilities.

For some of the channels, we can decay really, really slow, meaning that we can retain this long context information across a very long range.

And at the same time, for the other channels, we can quickly forget the information from the past indices to refresh it and observe new information.

And this is to increase the expressivity of this model.

And of course, to leverage modern GPUs, we have to use this Chang-Kwai's formulation so that we can parallelize the computation on modern GPUs.

So the first equation here is the Chang-Kwai's formulation of kimi-linear.

But as you can see, this is going to bring massive infrastructure challenges because of this newly introduced alpha term.

Because now it is a matrix instead of a scalar, it cannot easily be factored out.

So to achieve an efficient implementation, we rewrite the entire equation into the bottom three equations.

So we introduce this matrix inversion operation, as well as introducing the cumulative decay factor so that we can implement this entire thing in parallel without sacrificing any efficiency.

And more importantly, this is not an approximation.

It's an exact mathematically equivalent formulation so that we can achieve much efficient implementation without sacrificing any loss in terms of performance.

So it's going to be as efficient as previous linear attention variants, but at the same time, much more expressive.

So these are some of the results that we obtained using a fair comparison.

So on the left-hand side, we see the performance on two different types of tasks.

So MMU is a short context task.

So for short context tasks, Kimi-Linear achieved a better performance compared to MLA and GDN.

And at the same time, for longer context tasks, such as Ruler, Kimi-Linear is also better than the other variants while being much more efficient compared to MLA.

And when we scale the context transfer to, for example, 1 million tokens or even longer, it's going to be much more efficient compared to the baselines.

And this is also the first architecture that can outperforms full attention across the board, including short context tasks, long input tasks, and long output tasks.

So these are two key dimensions that we are interested in.

And the third dimension is the agent swarms.

So here is a diagram to showcase how we design this agent swarm paradigm to solve some of the more complex tasks compared to single agent paradigms.

So here, we have an orchestrator, or you can call it a main agent.

It's responsible for orchestrating tasks.

It has different options.

For example, we can spawn a group of sub-agents and assign new tasks to these sub-agents.

Or you can collect the results from the return of these sub-agents.

And you can perform this process in an iterative way.

And at the end of the day, you can accomplish a more complex task compared to using one single agent.

And it's analogous to human society.

For example, if we build a company, we need different roles.

And we need, for example, an orchestrator, or maybe we need a CEO, to decompose and assign the tasks to different roles.

And at the end of the day, the entire organization is going to have to move towards this same goal.

And here, for example, in this case, maybe you have the AI researchers, you have the web developers, you have physics researchers, and they can study different topics.

And at the end of the day, you just collect the results and spawn a group of fact-checkers, and web developers, and file downloaders to assemble the results into a single report.

And this is another perspective to look at this new paradigm.

So the x-axis is the complexity of the task.

And the y-axis is the execution time.

And the complexity of the task is measured by the accuracy of a group of models on such tasks.

So we can see, with agent swarms, it's going to substantially reduce the execution time compared to single agents.

It's going to be more effective.

And this means that we can scale this agent swarm paradigm to, for example, if you run this agent swarms with 100 or maybe even 1,000 agents, you can accomplish a complex task within a certain period of time that is tolerable for producing real economical value.

And we can certainly scale it in different dimensions.

We can scale the inputs.

For example, we can download and read hundreds of sources, or even maybe thousands of doses in parallel.

Or you can write a 100-page literature review in parallel.

Or you can take actions at scale.

You can perform data analysis for 10 different tasks.

And also, it is orchestration at scale.

You have to learn to design subtasks and aggregate the results.

And technically, we define some new objective functions to guide the learning process of our agent swarm system.

So there are three reward functions, reward objectives that are considered here compared to the conventional single agent RIO learning.

So the first term is what we call the instantiation reward.

It incentivizes sub-agent instantiation to prevent this serial collapse phenomenon from happening.

So basically, we don't want it to default to single agent execution.

We want to encourage the parallel executions, especially when it's early stage in training.

And of course, we can decay the weight for this instantiation reward term over training course, because when it learns parallel execution, we can reduce the weight.

And the second term here is finished reward.

And it is used because we observe one of the things in training that some of these subtasks are just created but never finished.

So it's almost like it's going to hack the first term by just spawning a bunch of sub-agents.

And the task might be too complex, or maybe the task just doesn't make sense.

And here, we use this finished reward to basically encourage that each of the subtasks should have a relatively high ratio of completion.

Instead of just spawning a bunch of pseudo-tasks, we need it to be meaningful.

So this is the second term that we use.

And of course, we use the same decay strategy.

We use a relative high weight at the beginning of training, and we decay it to a relatively low weight at the end of training.

And of course, the third term is the standard term.

It's the outcome reward.

It's going to measure whether the entire task is completed.

And then we're going to add these three terms in our reinforcement learning system.

And of course, we have to build the entire infrastructure, because right now, you need to support the parallel execution.

And then you'll need to support different reward functions and to maximize the efficiency of the entire agent swarm system.

So here are three different things that we have tried scaling.

The Mule and Clip Optimizer improves token efficiency.

And Kibme Delta Attention in the Kibme Linear Architecture improves long context.

And we also have the agent swarms paradigm to further create a new dimension of scaling.

And all of this put together, we created Kibme K2.5, a new model that we just released over one month ago.

Here's a short video to demonstrate some of its capabilities.

So let's take a look at Kibme K2.5., So yeah, there are a lot of interesting things, capabilities that we discover from the model.

For example, it merges the visual capabilities with coding capabilities.

So a lot of new things just emerge out of it.

It can read a video and then produce a website that sort of replicates or style transfer the original video.

And all of this are due to successful and stable training at the pre-training stage.

So this is also one of the most beautiful curves that I observed in my life.

So this is the training curve of the K2.5-based model.

So as you can see, it went through over 15 trillion tokens.

And of course, in K2.5, we additionally trained another 15 trillion tokens.

And the entire training process is just so stable.

There's no lost spike, especially when we introduced this new muon optimizer.

We didn't observe any spike.

And this smooth, stable training process produces a very stable outcome, a very strong-based model that we can fine-tune on top of it to achieve new capabilities, as we introduced and saw in the video.

And this is also, of course, trained on NVIDIA H800 GPUs.

And each node in this H800 cluster contains two TB RAM and eight GPUs.

They are connected by M-Link. And another key innovation of K2.5 is that it is the first open model with native joint vision text capabilities.

So if you look at previous open models, usually their visual capabilities are added on top of a text space, meaning that, for example, if you train the text models for 20 trillion tokens, and then on top of it, you do another 2 trillion sort of a post-training process to add additional visual capabilities on top of it.

But for K2.5, it's different in the sense that we fuse the training process of vision and text from day one.

So it's called early fusion here.

We start from 0% of the progress.

So from day one, we're going to merge the vision and text tokens.

And as shown in our preliminary experiments, it outperforms our late fusion.

And some of the new capabilities that we observe also come from this training recipe.

For example, if you want to do vision to code, you really have to merge vision and text into a single brand to achieve that.

If you separate these two brands, it's not going to happen.

You have to align these two modalities into a shared embedding space, a shared representation space, so as to achieve this.

And another interesting thing that we observe is that these two modalities can actually enhance each other.

So that's been long been a challenge, that if you add vision capabilities into a text model, it's going to somewhat hurt the text performance.

But here, we found that if you train it properly, these two modalities can actually enhance each other.

So this is one of the key findings that we observe in our training.

So first, vision improves text.

So this is so interesting.

So before vision RL, the performance in the first column.

And then we have the performance after vision RL.

So here, vision RL refers to a process that we only use vision tasks.

So there is no text task involved here.

We only have vision tasks.

For example, we teach the model how to count, how to answer some of this visual QA problems without any, for example, math, any coding problems in this space.

But we observe that it's going to improve the performance for even reasonably heavy text tasks.

And on the other hand, text also improves vision.

If you have a very strong text space, you actually don't need any vision SFT data in the training process.

And this is the approach that we adopt.

So it's called zero vision SFT.

Basically, we have basically zero vision SFT data.

And the only SFT data that we have is the text SFT data.

And then we do a joint RL over text and vision.

And you can see that we can achieve almost state of the art performance across the board on vision tasks without any vision data.

So it's clear that if you have a strong text space, it's also going to improve the vision if you align these two modalities into a shared space in your pre-training.

And also, these are some of the examples of, yeah.

As I was showing in the video, it demonstrates strong capabilities of visual design in front-end coding.

And this also emerges from our vision text for training.

So after all this, so this is all about KBK 2.5.

And as you probably know, we released our new architecture yesterday in our tech report.

It's called Attention Residue.

So here, I'm also going to briefly talk about our new work, which serves as a sneak peek into our next generation architecture that we're probably going to adopt in our later models.

So here, the motivation is quite simple.

Can we apply some of our techniques that we use in the temporal dimension?

And we just take some of the inspirations and then apply it to the depth dimension.

So and it starts from this residual connection.

So I still remember listening to Kami's talk at a tutorial in ICML 2016, 10 years ago.

So it was a brilliant idea.

So basically, before ResNet, nobody was able to train deep networks.

If you increase the depth, if you increase the number of layers for neural networks, nobody was able to train it.

Because you observe this gradient exposure, gradient vanishing, all this stability issues.

But then after the introduction of ResNet, we can train an arbitrarily large number of layers.

You can stack as many layers as you want.

And you don't have to worry about the training stability issue and stuff.

And as discussed in Ilya's talk two years ago, it basically says that residue connection is a variance of LSTM, but just rotated 90 degrees.

So how do you understand this?

If you look at LSTM, it's a variance of recurrent net.

And it's a recurrent model process.

So we're going to take the hidden states from the last step.

And then we're going to have some gating mechanism, some function to produce the current states.

And if you look at the depth dimension, residue connection is basically the same.

We're going to take the outputs from the last layer.

And then we're going to apply some sort of function on top of it to produce the current outputs of the current layer.

It's just the formulation is different.

For example, for residue connection, we're going to use a fixed addition.

We're going to have this.

We're going to add the previous hidden states with the current outputs.

It's just the formulation that's different.

But the basic idea is the same.

It's a recurrent net applied in the dimension of depth.

But on the other hand, we can think about reformulating this function.

Instead of having an LSTM, can we have an attention in the dimension of depth?

And it's going to create new possibilities, because attention has been demonstrated to be so successful in the transformer era.

So what we're going to do is not just to take the last hidden state, but we're going to consider all the previous hidden states and use the attention operation, the attention mechanism, to assemble and aggregate all of these previous hidden states to compute the current state.

So this is exactly attention rotated by 90 degrees.

It's sort of we view it as a natural generalization of residue connections in the LSTM analogy.

And here is the detailed formulation.

So on the left-hand side is a standard residue.

As I said, it's basically LSTM rotated by 90 degrees.

And the second figure is attention rotated by 90 degrees.

So what we do is to collect all the previous hidden states and have a simple attention operation on top of it to produce the current layer's outcome.

And of course, to increase the efficiency, to reduce the infrastructure, for example, communication and memory overhead, we also designed a new variant called block attention residue on the right-hand side.

So basically, the idea is also simple.

We're going to divide all the layers in the neural networks into multiple blocks.

For example, each block can contain, say, 16 layers, or it can contain maybe four layers.

And then for each block, we're going to apply this attention residue only on the outputs of each block.

But within each block, we still adopt this standard residue.

So this is going to reduce a lot of overhead while having minimal loss in terms of training accuracy.

And these are some of the impressive results that we achieved on this new architecture.

So on the scaling law, we can improve the token efficiency by 24%, meaning that if you have 50 trillion high-quality tokens, now you just magically have over 60 trillion tokens.

And then for the validation loss, you can also observe that it's consistently lower than the original curve, demonstrating this stability across optimization and also achieves the best improvement on some of this coding, math, and reasoning heavy tasks, as shown in the benchmark results of GPQA, math, and human eval.

So the entire community keeps moving forward.

And we're happy that we're able to contribute to the community with new technologies.

And some of these technologies have been sort of standard and de facto for a long time.

But as you can see, we still see a lot of opportunities to improve it, to have revolutionary new design to achieve better performance.

If we multiply all of these scans together, you can actually have a much better model.

So Adam was invented in 2014.

And now we scale an open-source Mule and Clip, a drop-in replacement for Adam.

And I'm sure that if you train a transformer LLM, it's going to be much better if you use Mule and Clip instead of Adam.

And attention was invented over eight years ago.

And then now we have KimiLinear, which is a linear version.

We don't have to use full attention across all layers.

We can have linear attention that performs better on short, long contexts at the same time.

And also, residue connections are now also challenged.

We scaled an open-source attention residue.

So I think one of the interesting things about our era is that we sort of adopt a different mindset for doing research.

So if we go back to 10 years ago, it's mostly about publishing a new idea.

But then I think the lack of the rigor of the experiments is very hard to produce solid experimental results.

But now we have this scaling ladder.

We have enough resources to train the model and run it at different scales.

We can have a whole set of benchmarks to measure the progress.

So it becomes easier to make confident and solid conclusion out of it.

And this is one of the reasons why we are observing new progress on this ancient techniques.

And I'm sure that we'll see more and more, especially in the open-source community.

I think we're going to have more and more even better architectural and optimization improvement in the next few years.

All right.

So to summarize, we're going to keep scaling our models.

And so these are three dimensions.

For example, we see different architectures and optimizers that optimize all three dimensions.

And we'll keep seeing new dimensions for scaling.

Agent swarms is not the end.

And we are glad that we can move forward with the entire open-source community to achieve better and better intelligence.

Thank you so much.