{
  "video_id": "Qf06XDYXCXI",
  "channel_slug": "statquest",
  "channel_handle": "statquest",
  "title": "Word Embedding in PyTorch + Lightning",
  "duration_seconds": 1922.0,
  "url": "https://www.youtube.com/watch?v=Qf06XDYXCXI",
  "upload_date": "",
  "transcript": "[Music]\nwood\nedding with po torch and lightning\nhooray stack\nQuest hello I'm Josh starmer and welcome\nto stack Quest today we're going to talk\nabout word embedding in pie torch plus\nlightning don't stress out about the\ncloud use\nlightning bam this stack Quest has also\nbrought to by the letters a b and c a\nalways b b c curious always B curious\nnote this stack Quest assumes you are\nalready familiar with word embedding if\nnot check out the quest also note you\ncan download all of the code in this\nstack quest for free the details are in\nthe pinned comment below in the stack\nQuest on word embedding we created a\nsimple neural network that converted the\nword or input tokens Troll 2 is great\nand Jim cata into numbers which we call\nword embeddings we also showed that\nthese word embeddings allow words that\nare used in similar contexts like Troll\n2 and Jim Kata to appear close to each\nother when we use the embedding values\nto plot each word on a graph\nbam now in this stack Quest we'll learn\nhow to build and train this simple word\nembedding network with P torch plus\nlightning first we'll do it from scratch\nusing just tensors and some basic math\nand then we'll simplify our code using\nthe pi torch linear function lastly\nwe'll learn how to use the pytorch\nembedding function to load and use\npre-trained word embeddings the first\nthing we do is import torch to create\nthe tensors we will use to store the raw\ndata and to provide a few helper\nfunctions then we import torch.nn to\ncreate the weights we will use in the\nnetwork and to bring in some other\nhelper\nfunctions then we import atom to fit the\nneural network to the data with back\npropagation then we import uniform to\ninitialize the weights in the\nnetwork and to give us the tools to\ncreate a large scale embedding network\nwith lots of training data we'll import\ntensor data set and data loader from\ntorch. ts. data now we import lightning\nas L to make it way easier to write our\ncode and for automatic code optimization\nand scaling in the cloud lastly we\nimport pandas matte plot lib and Seaborn\nso that we can draw some pretty\ngraphs now we will create word\nembeddings for these two simple\nsentences Troll 2 is great and Jim Kata\nis great so we will need to create an\ninput for for each unique token in the\ntraining data and these inputs will\neventually connect to this word\nembedding\nNetwork now if we wanted to run Troll 2\nthrough the network we would put a one\nin the input for Troll 2 and a zero in\nall of the other inputs oh no it's the\ndreaded terminology alert when we\nspecify the inputs like this where one\ninput gets a one and everything else\ngets a zero\nit's called one hot\nencoding so if we want to run the word\nis through the network we specify the\none hot en coding with a one for is and\na zero for everything else likewise this\nis the one hot en coding for great and\nthis is the one hot en coding for Jim\ncata now going back to the one hot en\ncoding for trol 2 in pi torch we can\ncreate a Four Element list with one in\nthe the first position for Troll 2 and\nzeros in all of the other positions to\nrecreate the one hot encoding for Troll\n2 likewise we can create similar one hot\nencoding lists for is great and Jim Kata\nand once we have a list for each\npossible input we make a list of these\nlists and because we're using pi torch\nwe convert the input lists into tensors\nwith torch. tensor and save the tensors\nin a variable called\ninputs now the goal is to create a\nsimple Network that can predict the\ntoken that follows a specific input for\nexample if the input token is Troll 2\nthen we want to predict the word is and\nthat means we want the output for is to\nbe one and the outputs for all of the\nother tokens to be zero in pi torch that\nmeans we want to predict the one hot\nencoding for\nis likewise when the input is\nis then we want to predict great and\nthus we want to predict the one hot\nencoding for great great is at the end\nof each phrase and nothing comes after\nit so in theory it shouldn't predict\nanything but if we had a larger training\ndata set it probably would predict\nsomething so in this example will just\npretend that it predicts Jim Kata lastly\nJim Kata just like Troll 2 predicts is\nnow let's combine the output lists in a\nlist and convert the output lists into\ntensors and save them in a variable\ncalled labels because in machine\nlearning that is what we call the known\nor ideal output values and now we are\ndone encoding the training data now that\nwe have the inputs and label\nwe can combine them into a tensor data\nset that we'll call data set and then\nuse data set to create a data loader\ncalled Data loader data loaders are\nsuper useful when we have a lot of data\nbecause one they make it easy to access\nthe data in batches two they make it\neasy to shuffle the data each Epoch and\nthree they make it easy to use a\nrelatively small fraction of the data if\nwe want to do a quick dirty training for\ndebugging but Josh we don't have a lot\nof data why are we using a data loader\nyou're right Squatch we don't have a lot\nof data but in a more realistic setting\nwe would so we might as well do it now\nokay anyway now that we have the data\ntaken care of let's write the code for\nthis simple word embedding Network when\nwe create a neural network in pi torch\nwe always start by defining a new new\nclass now because we're coding our first\nword embedding network from scratch\nwe'll call this word embedding from\nscratch and in order to make coding\nsuper easy We'll Inherit from lightning\nmodule then just like we always do we\ncreate an initialization method for the\nnew class this method will create and\ninitialize all of the weight tensors\nthat we need to implement the embedding\nNetwork and it will also create the LW\nfunction that we'll use during\ntraining then we'll create a method\ncalled forward that makes a forward pass\nthrough the embedding Network then we'll\ncreate a method to configure the atom\nOptimizer and lastly we'll create a\nmethod called training step to calculate\nthe loss which in this case will be the\ncross entropy loss the cross entropy\nloss function will quantify the\ndifference between what we want the out\noutput to be and what we actually get\nfor the output note there's a lot more\nto be said about the cross entropy loss\nfunction so if you're curious check out\nthe quest now let's start by coding the\nanit method the first thing we do is\ncall the initialization method for the\nparent class lightning module this is\nsimply required whenever we inherit from\na class in Python in this case this will\nallow us to take advantage of all the\nfeatures that lightning\noffers now we need to create and\ninitialize the weights for the network\nand we're going to do this by using a\nuniform distribution to randomly select\nan initialization value for each weight\nspecifically we're going to use this\nuniform distribution that goes from 0.5\nto\n0.5 to Generate random numbers for the\nweights the shape of this distribution\nshows the that all values between 0.5\nand 0.5 have the same likelihood of\ngetting randomly selected hey Josh why\nare we using values between 0.5 and 0.5\nto initialize the weights good question\nSquatch we're using this specific range\nof values in order to match up with what\nwe will do in the second part of this\ntutorial when we use the pi torch linear\nfunction to do the math and the L linear\nfunction selects a range of values based\non the number of\ninputs okay anyway in order to use this\nuniform distribution to Generate random\nnumbers for us we create a variable\ncalled Min value and set it to\n0.5 the minimum value we want to\nrandomly select and we create a variable\ncalled max value and set it to 0.5 the\nmaximum value we want to randomly select\nnow we create a parameter for the first\nweight associated with the first input\nand use uniform. sample which we\nimported earlier to initialize it with a\nrandom number when I did this the first\nweight associated with Troll 2 was\nrandomly set to\n0.38 then the second weight associated\nwith Troll 2 was randomly set to\n0.42 now we just create an initialize\nparameters for all of the other weights\nassociated with each input bam now we've\ncreated and initialized all of the\nweights associated with the inputs\nlikewise we create and initialize\nparameters for all of the weights\nassociated with each\noutput now we have all of the weights\ninitialized for our word embedding\nNetwork\nbam now the last thing we need to do in\nour anit method is give our class access\naccess to the Cross entropy loss\nfunction and we do this by calling nn.\ncross entropy loss and saving it in a\nvariable called\nloss now that we're done coding the anit\nmethod we can use those weights to code\nthe forward method to make a forward\npass through the embedding Network for\nthe forward method the input is a list\nthat contains the one hot encoding for\none of the input tokens for example the\ninput might be the one hot encoding for\nTroll 2 however when it is passed to the\nforward method it comes wrapped up in an\nextra set of\nbrackets so the first thing we do is\nremove those brackets by setting input\nto be the first\nelement now we multiply each input value\nby its corresponding weight that goes to\nthe activation function on top and add\nthe products together and we save that s\nin a variable called inputs to top\nhidden then we multiply the inputs by\nthe weights to the activation function\non the bottom and add the products\ntogether and save the sum in a variable\ncalled inputs to bottom hidden so now we\nhave the code for the first part of the\nword embedding Network and since the\nactivation functions are identity\nfunctions which means the input is the\nsame as the output we can multiply\ninputs to top hidden and inputs to\nbottom hidden by the next set of Weights\ndirectly for example we can multiply\ninputs to top hidden by the first weight\ngoing to the top output and then\nmultiply inputs to bottom hidden by the\nsecond weight going to the top output\nand then add those two products together\nand save the result in a variable called\noutput one then we do the same same\nthing for the other\noutputs bam now we've done all the math\nup to the softmax function and that\nactually means we are done making a\nforward pass through the embedding\nNetwork because the loss function that\nwe're using for back propagation nn.\ncrossentropy loss does the soft Max for\nus so the last thing we need to do is\npackage up the output values using\ntorch. stack and save everything in a\nvariable called output pre- soft Max\nnote if instead of using torch. stack we\njust returned a list of the output\nvalues by wrapping them up in square\nbrackets then the gradients would get\nstripped off and we would not be able to\ndo back\npropagation so by using torch. stack we\ncan return a list that preserves the\ngradients anyway the last thing we do in\nthe forward method is return output pre-\nsoft\nMax now that we have the forward method\nwe are ready to configure the\noptimizer and configuring the optimizer\nin this case atom is so easy we can just\nreplace the pseudo code with the real\ncode we pass atom the parameters we want\nto\noptimize and we set the learning rate LR\nto\n0.1 hey Josh why did you set the\nlearning rate to\n0.1 because our example is pretty simple\nand I wanted to train relatively quickly\nI tested out a relatively large learning\nrate 0.1 and it worked okay now let's\ntalk about the training step method\nwhich we'll use to calculate the loss\nthe training step method takes a batch\nof training data and the index for that\nbatch and the first thing we do is split\nthe batch of training data into the\ninput\nand the labels which are the ideal\noutput\nvalues then we run the input through the\nnetwork up to the softmax function by\npassing it to the forward method for\nexample if we run the one hot encoding\nfor Troll 2 through the untrained\nNetwork these are the values that the\nforward method will return we then run\nthose values along with the ideal values\nthrough the loss function nn. cross loss\nthen runs the output values through a\nsoft Max function and quantifies the\ndifference between the soft Max output\nin the ideal values and we save that\ndifference in a variable called\nloss and then return the loss now at\nlong last we've made it through all the\ncode needed to create word embeddings\nfrom\nscratch we create and initialize the\nweight tensors and create the loss\nfunction in the init method we make a\nforward pass through the embedding\nnetwork with the forward method\nconfigure the atom Optimizer with\nconfigure uncore\noptimizers and last but not least\ncalculate the loss with training step\nbam now let's use the new class we just\nwrote to create a new word embedding\nNetwork that we'll call model from\nscratch and let's print out the randomly\nselected weight values that it starts\nout with here we just have a for Loop\nthat iterates over all the named\nparameters in the network and for each\nparameter it prints out its name and\nvalue and here's the\noutput now to be honest this list of\nnumbers is kind of hard to read so let's\norganize it into an easyto read data\nframe so the first thing we do is put\nthe weight values into a dictionary the\nfirst first part label W1 contains the\nweight values for each input that goes\nto the activation function on top note\nwe're using the item method to get the\nweights because it Returns the tensor\nvalues as python\nnumbers the second part label W2\ncontains the weight values for each\ninput that goes to the activation\nfunction on the bottom then we just\nlabel the tokens and inputs and save the\ndictionary in the very aable called data\nand then transform data into a pandis\ndata frame called\nDF and print out\nDF and this is what the data frame looks\nlike now we can easily see that the\nweights for Troll 2 and Jim Kata are\nrelatively different even though they\nboth represent movie titles that are\nused in the same\ncontext this table is pretty helpful for\nmaking the embedding values easy to look\nat but a graph would make them even\neasier to look at this graph has the\nweight values to the top activation\nfunction W1 on the x axis and the weight\nvalues to the bottom activation function\nW2 on the Y\nAIS with a graph it's super easy to see\nthat the embedding values for troll two\nare very different from the values for\nJim Kata to create the graph the the\nfirst thing we do is call the caborn\nfunction scatter plot and we pass\nscatter plot the data frame DF that we\njust\ncreated and we tell scatter plot that we\nwant to use the weights that go to the\ntop activation function W1 on the X\nAIS and the weights that go to the\nbottom activation function W2 on the Y\nAIS now if all we did was call scatter\nplot then we'd end up with this scatter\nplot and while this scatter plot is\nsuper cool it would be much cooler if\neach dot were labeled with the word or\ntoken that it represented like this so\nin order to add the tokens as labels to\neach point we call the Matt plot lib\ntext function and we pass in the X and Y\nAIS coordinates for the point in the\nfirst row in the data frame and the\nvalue for the token then then we specify\nhow we want the text aligned the font\nsize the font color and lastly the font\nweight then we do the same thing for\neach row in the data frame and then we\ncall PLT\ndoow and we get this super cool looking\nscatter plot like we mentioned earlier\nwe can now see that the embedding values\nfor Troll 2 and Jim cata are pretty\ndifferent and that means we need to\ntrain our embedding\nNetwork we start training by creating a\nlightning trainer called trainer and\ntell it to train for at most 100 Epoch\nwhich means we will do back propagation\nfor every weight using the training data\nat most 100 times now we call the\ntrainers fit method and pass it the\nembedding Network called model from\nscratch and the training data called\nData loader\nin theory it should only take a few\nseconds to train our simple embedding\nNetwork and when it's done we recreate\nthe data frame that has the weights or\nembedding values for each token and we\ncan either stare at the embedding values\nin the data frame or we can draw a\nscatter plot of the tokens just like\nbefore note because the labels for Troll\n2 and Jim cata are overlapping and hard\nto read I then added little offsets to\nwhere Troll 2 and Jim cata were printed\nand now we can see that after training\nthe embedding Network the embedding\nvalues for Troll 2 and Jim cata are very\nsimilar which is great since they are\nused in similar contexts\nbam now that we have trained our\nembedding Network we can see what it\npredicts when we use Troll 2 as the\ninput remember from when we created the\ntraining data that we want Troll 2 to\npredict\nis so the first thing we need to do is\ncreate a soft Max function because we\ndidn't have to explicitly use it in our\nmodel note we set dim equal to zero so\nthat we can apply it to rows of output\nvalues if we set dim equal to one then\nwe would apply it to Columns of values\nnow we pass Troll 2 as a one hot encoded\ntensor into model from\nscratch and we run the output values\nthrough the the soft Max and then round\nthe output of the soft Max to two\ndecimal\nplaces and finally print out the result\nand we get the one hot encoded tensor\nfor is which is correct\nbam likewise we can verify that all of\nthe other inputs to our embedding\nNetwork create the correct\noutputs okay now that we know how to\ncreate and train a simple word embedding\nNetwork from\nscratch let's make our lives a little\neasier by using the P torch linear\nfunction to create the same network so\nlet's create a new class called word\nembedding with\nlinear and again in order to make\ntraining super easy will inherit from\nlightning\nmodule then create the anit method that\nwe use to create and initialize the\nweights and like always we'll call the\nanit method from the parent class\nnow comes the interesting part instead\nof calling nn. parameter to create and\ninitialize each weight in the network we\nonly have to make two calls to nn.\nlinear the first call creates the\nweights between the inputs and the\nhidden layer in features equals 4 means\nwe are connecting four inputs to two\nnodes specified with features equals 2\nin the hidden layer\nin other words this call to nn. linear\nwill make four weights for each of the\ntwo nodes in the hidden\nlayer and since we don't need any bias\nterms we set bias equal to\nfalse the second call to nn. linear\ncreates the weights between the hidden\nlayer and the\noutputs it creates two weights within\nfeatures equal to two for each of the\nfour outputs without features equal to 4\nand again since we don't need any bias\nterms we set bias equal to\nfalse now the last thing we need to do\nin our anit method is give our class\naccess to the Cross entropy loss\nfunction now we need to code the forward\nmethod that makes a forward pass through\nthe\nnetwork the cool thing is that all we\nhave to do to calculate the sums before\nthe activation\nfunctions is pass the input to the\nlinear object input to Hidden that we\ncreated in the anit method and save the\nsums in a variable called hidden the\nlinear object input to Hidden does all\nof the multiplication and addition for\nus\nbam note now that we are using linear to\ndo the math we no longer have to strip\noff the extra brackets from input like\nwe did\nbefore anyway because the input to these\nactivation functions is the the same as\nthe output we can just ignore them and\npass hidden to the second linear object\nwe created hidden to Output hidden to\nOutput calculates the output sums from\nthe activation functions and we save\nthose output values in output\nvalues now remember that we don't need\nto calculate the soft Max because the\nloss function nn. crossentropy loss does\nit for us so all we have to do is return\nthe output values\nbam the next thing we do is create the\nconfigure optimizers method and just\nlike before we use the atom Optimizer\nand pass it the parameters we want to\noptimize and set the learning rate to\n0.1 the training step method is also\nsimilar to what we did\nearlier and that means that this\nembedding network is contained in this\nclass definition by using inn. linear we\nsignificantly reduced the amount of code\nwe need compared to when we did\neverything from scratch we can see the\ndifference when we shrink the original\ncode for the word embedding from scratch\nclass down so that we can fit it on a\nsingle screen and compare it to the size\nof the word embedding with linear class\non the same hardto read\nscale\nbam now we can create a new model model\nlinear with our new class word embedding\nwith\nlinear and just like before put the\npre-trained word embedding values into a\ndata frame called\nDF note because we used nn. linear to\ncreate the weights we access them with\nweight and we call detach to remove the\ngradient from the\ntensors and we use zero and one to index\nthe weights that go to the top and\nbottom activation\nfunctions and lastly convert the tensor\nto a num Pi array with num Pi now when\nwe print out our data frame\nDF we get this nicely formatted table\nand we can draw a scatter plot of the\ntokens just like before and that gives\nus a scatter plot that looks like this\nand the graph suggests that the\nembedding values or weights are not yet\noptimal because Troll 2 and Jim are so\nfar from each other so just like before\nwe can train the model for 100\nEpoch and after training we end up with\nthese\nweights and the weights going from the\ninputs to the hidden layer are the new\nembedding\nvalues and when we redraw the scatter\nplot with the new embedding values we\nsee that the values for Troll 2 and Jim\ncata are similar double\nbam now that we we know how to create\nword embedding networks from scratch and\nwithin in.\nlinear and we can create word embeddings\nthat put words and tokens used in\nsimilar context near each other let's\nlearn how we can load and use\npre-trained word embedding values with\nnn.\nedding Note in this example we're just\ngoing to load and access the embedding\nvalues we\ncreated but in practice we might want to\nload and use the word Tove embedding\nvalues which have 100 values per token\nand millions of\ntokens or the embedding values created\nby a Transformer like chat\nGPT however before we get started let's\njust print out the embedding values from\nthe last Model we trained model\nlinear we access the embedding values in\nmodel linear just like we did when we\ncreated the data frame except now we\ndon't have to worry about the gradient\nan and we get two lists of embedding\nvalues specifically the weights are\narranged in two\nrows the first row corresponds to the\nweights that go to the top activation\nfunction and the second row corresponds\nto the weights that go to the bottom\nactivation\nfunction now the problem with this is\nthat nn. embedding expects the weights\nto be in columns just like in the data\nframes we created the good news is that\nconverting rows into columns is super\neasy so with that said let's load and\nuse these pre-trained word embedding\nvalues with nn.\nedding we start by creating an nn.\nembedding object and we pass the\npre-trained weights in with from\npre-trained and we use T to transpose\nthe rows of Weights into\ncolumns lastly we save the the new\nembedding object in word\nembeddings we can then verify that we\ndid things correctly by printing out the\nweights and we see that the weights are\nnow arranged in two\ncolumns now we can print out the\nembedding values for the first input\nTroll 2 by passing in a tensor with the\nfirst index value zero and the embedding\nvalues match what we expect so we know\nwe did things\ncorrectly accessing the embedding values\nby index is fine but we can also make\nour lives Easier by creating a\ndictionary that maps The Tokens to their\nindices and now we can more easily\naccess the embeddings with the token\nitself rather than the\nindex and that's all there is to loading\nand accessing pre-trained weights into\nan nn. embedding object we can now use\nour embedding object word embeddings and\nconnect it to a larger neural network\nlike a Transformer also before we go I\njust want to remind you that you don't\nhave to type this code yourself instead\nyou can download it and I wrote tons of\ncomments that explain every little\ndetail just like this stack Quest the\nlink is an app pinned comment\nbelow triple\nbam now it's time for some Shameless\nself-promotion if you want to review\nstatistics and machine learning offline\ncheck out the stack Quest PDF study\nguides and my book the stat Quest\nIllustrated guide to machine learning at\nstat quest.org there's something for\neveryone hooray we've made it to the end\nof another exciting stack Quest if you\nlike this stack Quest and want to see\nmore please subscribe and if you want to\nsupport stack Quest consider\ncontributing to my patreon campaign\nbecoming a channel member buying one or\ntwo of my original songs or a t-shirt or\na hoodie or just donate the links are in\nthe description below all right until\nnext time Quest on",
  "transcript_chars": 24784,
  "ingested_at": "2026-05-15T10:55:05.762274+00:00",
  "source": "channel",
  "yt_meta": {
    "view_count": 51706,
    "like_count": 1015,
    "channel_id": "UCtYLUTtgS3k1Fg4y5tAhLbw",
    "categories": [
      "Education"
    ],
    "tags": [
      "Josh Starmer",
      "StatQuest",
      "Machine Learning",
      "Statistics",
      "Data Science"
    ]
  }
}