lucidrains implementation
starting with the lucidrains(goat) implementation see what improvments we can make from their own details about it are
- it trains the transformer along it
- it only adds the context at the starting as in the paper
- 2 layer mlp
- persistent memory i can probably just add more this config and get better results
why does his script not log the val_loss on wandb :sob:
anyways so far the val loss looks simliar to training im gonna wait for it to fall below 0.3 and see from there
aight so ummmm the lucidrains implentation works fine really but it builds the transformer along with it i don’t really want that my goal is to have like an extension that i can just stick to any llm the extension should be able to learn how to distill the KV pairs of the LLM into the LLM itself and retrieve at the right moment,, i don’t see the need for persistent memory either… so here’s the goal
- a model that sticks at the top in first (few) layer and does what it is supposed to let’s call it “titans at head”
- a model that joins after each layer adds relevant key pairs after FFNs let’s call it “titans stream” now each of them will have their own configs and params etc to actually compare both of these and what performs the best i want to take the parameters equally in both the cases. so let’s say 2 layers each layer after MLP in titans stream so the model i want with titans at head should be proportional to that but denser perhaps N/2 where N is the parameters of the LLM i can also try giving the titan layer every i number of layers in the LLM ig in titans stream the the good thing about titans head is that it can just be stuck to any LLM and it should learn to learn to distill KV pairs and change its states every-time it is initialized but if i pull off custom titans stream that would be crazy so let’s see
let’s start with titans head
evals
need to build an eval first of all to compare against the normal LLM and my extension for titans hmmmm.
so gemma 3 4b has around the default context of 4096 in cache_length tokens with caching and 1024 in the sliding attention. we can increase the cache length to upto 128k tokens but i’m not sure if my system can handle that but gotta try
sooo what kind of tests dod we want??
OKAY fuck what do we do for tests
- have a big pdf ask multiple choice questions
- specific page specific line??
- semantic content that couldn’t be in the embedding so like code snippet is the query and it finds the latex functions
- SAT style problems
- specific citation of multiple documents sentence wise
- detective style problem with cross linking facts
ideas
- add kblam square attention for incorporating in context
titans head
this is what the modular code looks like its very complex
graph TB subgraph Input X[Input Sequence X_t] end subgraph Memory_Module WK[Key Projection W_K] WV[Value Projection W_V] MN[Memory Network M] MU[Memory Update] MS[Momentum State S_t] end subgraph Segment_Processing WQ[Query Projection W_Q] Cat[Concatenation] MHA[Multi-Head Attention<br/>8 heads] Out[Output Projection] end Memory Module Flow WK --> |K_t| MN WV --> |V_t| MN MN --> |h_t| Cat MN --> MU MU --> |Update| MN MS --> |Momentum| MU Add styles classDef module fill:#f9f,stroke:#333,stroke-width:2px classDef process fill:#bbf,stroke:#333,stroke-width:2px class Memory_Module,Segment_Processing module class WK,WV,WQ,MN,MHA,Out process
i want to go simple i have one more idea
- have a vector db compare the output vector find the closest match and pass those
i initially only want to try to convert h_t without the whole attention module in between but maybe in later experiments i can add them to see if it gives any actual improvment also want to do two different tests with h_t raw conversion versus look up conversion
basic run
now obviously there are a lot of things to fix but the prime i realized was that the loss needs to be different
since the output is a vector we can’t define the number of tokens since it is sentence embeddings so gotta do it someway to MSE loss or something, also i can’t quite understand what exactly is happening in the code since i am a fucking retard
there’s a lot of mismatch so
we might need to change the eq params to be data adaptable
k_t for memory train
did another run this time k_t for comparing to train the memory module
i also updated the token length stuff to vary from 2k to 128k tokens
this is bad because it overfit and was spouting random bs
this is a hard probelm to solve 
also both training and infer are embedding stuff at sentence levels the fucking idiot claude doesn’t get it gotta fix it
new dataset let’s go
so since the dataset is context and pairs of quesiton and answer i am still doubtfull on what i should do i can either go token by token since this is a pretraining but since we are using sentence embedding we can also just use it to predict the next sentence, my reasoning is i think i should just have the context and the question as seperate embeddings and get it to predict the answering embedding it’s essentially retrievel but if that were possible it would just be a fucking general intelligence we need it to learn that the embedding inside the question is the answer
we need skip connections
we gotta do a lot of tests today coz from tomorrow we need to start
new prompt to discuiss with chatgpt/claude
i am trying to make a new model
it is based off of titans architecture the google paper specifically the MAC arch
what we want is given the context and a question we get the relevent info that's it. now i have downloaded the MS marco dataset and am using SONAR embeddings since the final output is an embedding. i am here to discuss with you what i should do for making it work like should i
a. make it token level
b. make it sentence/set level
c. something you suggest
token level is i embed each token in the context and the question i concatenate them and get the output and do the loss over it
entire sentence or set is if i embed the entire context and then concatenate the entire question dataset (since it is SONAR and sentence level embedding) and have it predict the answer tokens/sentence and do loss over it
we can do a normal MSE loss or some other type of loss cross entropy or cosine similarity with a mix of mse or something
think into consideration that we want this mode to be lightweight like how the mac architecture is just two layers in the paper and we want it to be such that it just gets the relevant tokens with respect to the sentence (we might need to edit the dataset for that)
think thoroughly give reasoning for everything and you are also allowed to dispute anything i have said or to call my idea out if it is not making any sense and i should'nt pursue it
ok so we do sentence level embeddings in the entire paragraphs
ok fuck me ms marco dataset idea was stupid, the whole point of the inner algo is it updates with each new token and is suppsoed to spread out the attention of the tokens that are forgotten or need to be taken in context i need to read the paper more carefully anyways
so the new or old strategy is to have this extension input the prev context and output some embedding and then have the model predict the next token and use that loss to work off of