Paper: FIITED: Fine-grained embedding dimension optimization during training for recommender systems (ICLR-2024)

Reviewer comment (Novelty):
The paper's motivation needs to be more explicitly articulated. The primary concern is the rationale behind the compression of embedding dimensions for standard recommender systems. When general recommendation models typically deploy on powerful cloud servers, the necessity to compromise model accuracy for size reduction seems unclear.

Author response:
Although recommender models are usually trained on powerful cloud servers, as the model size keeps increasing, memory still tends to be the bottleneck. Also, as training machine learning models on edge devices becomes more and more popular nowadays due to privacy concerns etc., memory resources may be fairly scarse in some use cases.

Model compression during training is a solution to this problem, and it does not necessarily lead to model quality deterioration. As briefly mentioned in Section 4.2.1, FIITED managed to achieve a 0.1% accuracy gain on the Criteo Terabyte dataset under a pruning rate of 93.75%. In other cases, FIITED can achieve superior compression rates while maintaining the model quality. Such results suggest that FIITED is able to find suitable dimensions for the embeddings while saving a considerable amount of training resources.

Reviewer comment (Experiments):
The selection of baselines for comparison appears limited. Other potential methods such as DNIS [4], PEP [6], AMTL [7], and AutoDim [5] could have been considered or at least discussed in terms of baseline selection. A clear criterion for selection would lend more weight to the experimental results. Merely mentioning ESAPN as the state-of-the-art may not suffice.

Author response:
Because a main goal of FIITED is to save memory during training, we have selected baselines that can do the same. AdaEmbed is one such example where entire embedding rows are pruned during training. ESAPN [2] is another example where embedding dimensions are adjusted during training, and it claimed that memory saving can be achieved by always storing embeddings at their current lengths, although no system design was provided to show how the fragmented free memory can be utilized. Nevertheless, it was included in the comparison.

Other potential methods are discussed in Section 2, including [3,5,6,7]. It will be interesting to see whether FIITED can reduce the model size more than these methods, but a crucial reason why they are not included in the baselines is that none of them [1,3-7] can actually save memory during training, and many of them increase the memory usage in training by introducing new structures to the recommender model. Although PEP [6] reduces memory usage in the retraining step, the pruning step still needs to store the entire unpruned embeddings as well as additional parameters used for pruning while training the model.

We will add the discussion of baseline selection and include citations [1,4] into the manuscript.

Reviewer comment (Reproducibility):
There are missing technical details that might be crucial for a thorough understanding and replication of the method. The division of the embedding table into m*k embedding blocks is mentioned, but the determination of 'm' and the grouping strategy for different feature values are not elaborated upon. A comprehensive study or discussion on the hyperparameters, especially concerning 'm' and 'k', would be valuable. As of now, only Figure 6 sheds light on the varying training times with 'k'.

Author response:
The embedding table is not divided into blocks. Every row in the embedding table (i.e., each embedding vector) is divided into K equal-sized chunks. For example, if the default embedding dimension is 128 and K=4, then the 4 chunks will be [0, 32), [32, 64), [64, 96) and [96, 128). Currently there is no grouping of feature values. Grouping was used in AdaEmbed to handle cases where there are different default embedding dimensions (e.g., one sparse feature may have a default dimension of 128 while another has 64). But this is not necessary for us, since in our experiments, there is only 1 default embedding size.

I re-read the submission and could not find the notion of blocks. The notation m was used in Algorithm 1 as the sample size when sorting the utility values. Please let us know where "m*k blocks" was mentioned so that we can improve our wording and presentation.

Reviewer comment (Experiments):
On what grounds were the baselines selected for comparison? Why were certain other potential methods excluded from the baseline?

Author response:
Because a main goal of FIITED is to save memory during training, we have selected baselines that can do the same. AdaEmbed is one such example where entire embedding rows are pruned during training. ESAPN [2] is another example where embedding dimensions are adjusted during training, and it claimed that memory saving can be achieved by always storing embeddings at their current lengths, although no system design was provided to show how the fragmented free memory can be utilized. Nevertheless, it was included in the comparison.

Other potential methods are discussed in Section 2, including [3,5,6,7]. It will be interesting to see whether FIITED can reduce the model size more than these methods, but a crucial reason why they are not included in the baselines is that none of them [1,3-7] can actually save memory during training, and many of them increase the memory usage in training by introducing new structures to the recommender model. Although PEP [6] reduces memory usage in the retraining step, the pruning step still needs to store the entire unpruned embeddings as well as additional parameters used for pruning while training the model.

We will add the discussion of baseline selection and include citations [1,4] into the manuscript.

Reviewer comment (Reproducibility):
In Section 3.2.1, you mention that individual pruning ratios for different chunks may be bigger or smaller than the global pruning ratio, and may change over time, depending on chunk utility values computed at runtime. Can you provide more details on how these utility values are computed and how they are used to determine the pruning ratios for each chunk?

Author response:
**Q2**. In Section 3.2.1, you mention that individual pruning ratios for different chunks may be bigger or smaller than the global pruning ratio, and may change over time, depending on chunk utility values computed at runtime. Can you provide more details on how these utility values are computed and how they are used to determine the pruning ratios for each chunk?

As mentioned in Section 3.2 and also in Q1, the utility value of a chunk in the i-th iteration is computed as follows:

u(i) = γu(i − 1) + a(i)g(i)

where a(i) is the number of accesses to the chunk in iteration i, g(i) is the L2 norm of the gradients computed for the chunk in iteration i, u(i − 1) is the history utility value of the chunk computed in the previous iteration, and γ ∈ (0, 1) is a decay parameter that reduces the influence of the history utility value. The initial utility value is 0 at the start of training for every chunk. The utility value is updated in this way in every iteration for every chunk. More details for how utility values are computed are also included in Q1.

As for deciding the pruning threshold, there are two ways. One is to manually specify pruning rates for each chunk. For example, below we show a simplified case of 4 embeddings with 2 chunks each. If we set the pruning rate to 50% for chunk 0 and 75% for chunk 1, then 2 chunks remain in the first chunk column, and 1 chunk remains in the second chunk column.

| | chunk 0 utility | chunk 1 utility |
|---|---|---|
| Embedding 0 | 1.3 (keep) | 1.7 (keep) |
| Embedding 1 | 0.9 (prune) | 1.6 (prune) |
| Embedding 2 | 0.3 (prune) | 0.2 (prune) |
| Embedding 3 | 2.1 (keep) | 1.5 (prune) |
| Pruning rate | 50% | 75% |
| Pruning threshold | <1.3 is pruned | <1.7 is pruned |

If manual specification is used, then the pruning ratio for a chunk column remains fixed during training.

The other way is to set a global pruning rate. For example, if a global pruning rate of 50% is used, then across all chunk columns, the 4 chunks with the highest utility values will remain, as shown below.

| | chunk 0 utility | chunk 1 utility |
|---|---|---|
| Embedding 0 | 1.3 (prune) | 1.7 (keep) |
| Embedding 1 | 0.9 (prune) | 1.6 (keep) |
| Embedding 2 | 0.3 (prune) | 0.2 (prune) |
| Embedding 3 | 2.1 (keep) | 1.5 (keep) |
| Pruning rate (global) | 50% | 50% |
| Pruning rate (chunk) | 75% | 25% |
| Pruning threshold | <1.5 is pruned | <1.5 is pruned |

If a global pruning rate is used, then the pruning ratio for a chunk column can fluctuate during training, because it is affected by utilities in other chunks.

Reviewer comment (Experiments):
In Section 3.3, you introduce the Virtually Hashed Physically Indexed (VHPI) embedding table design, which is a modification of AdaEmbed. Can you provide more details on how VHPI mitigates the issue of memory fragmentation and how it compares to other embedding table designs?

Author response:
**Q3**. In Section 3.3, you introduce the Virtually Hashed Physically Indexed (VHPI) embedding table design, which is a modification of AdaEmbed. Can you provide more details on how VHPI mitigates the issue of memory fragmentation and how it compares to other embedding table designs?

Compared to a traditional 2D array design, the following example shows how FIITED solves the memory fragmentation issue. Consider a case of 4 embeddings, each having a different dimension with a maximum dimension of 8. Here we use numbers to indicate unpruned parts and _ for pruned parts.

Embedding 0: 0000____

Embedding 1: 111111__

Embedding 2: 22_____ 

Embedding 3: 33333333 

In the traditional design, the memory will look like: 0000____111111__22_____33333333. The free memory is scattered around in separate small segments, and it is not easy to utilize.

In FIITED's design, only unpruned chunks are stored. In the VHPI hash table, addresses of the unpruned chunks are kept in order to read embedding chunks from the VHPI embedding table. For example, the VHPI hash table and the embedding table can look like this for the case above (with K=4, K is the number of chunks):

--------

**VHPI hash table**:（each entry has addresses of chunks, chunk utility values and a bit mask showing whether the chunks are pruned; invalid chunk addresses for pruned chunks are shown as N/A for clarity)

Entry for Embedding 0:     (9, 5, N/A, N/A, chunk utility values, 1100)

Entry for Embedding 1:     (4, 6, 1, N/A, chunk utility values, 1110)

Entry for Embedding 2:     (0, N/A, N/A, chunk utility values, 1000)

Entry for Embedding 3:     (2, 8, 3, 7, chunk utility values, 1111)

--------

**VHPI embedding table**:

index 0,1 :     22,11

index 2,3 :     33,33

index 4,5 :     11,00

index 6,7 :     11,33

index 8,9 :     33,00

--------

For example, if we want to access embedding 0, we will first check its hash table entry, and see that there are 2 unpruned chunks at index 9 and 5 in the embedding table, and 2 pruned chunks. Then, we can get the embedding chunks from the VHPI embedding table by reading the values at index 9 and 5. The chunks are then concatenated to construct embedding 0.

Now the memory looks like 22113333110011333300, with unpruned chunks stored together and free memory not in segments. 

Most recommender systems nowadays use the traditional 2D array storage structure. The VHPI storage system proposed by FIITED enables real memory saving during training.

Reviewer comment (Reproducibility):
In Section 4.2, you evaluate the performance of FIITED on four datasets, including three click-through rate prediction datasets and one classification dataset. Can you provide more details on the hyperparameters used in your experiments, such as the learning rate and batch size, and how they were selected?

Author response:
**Q4**. In Section 4.2, you evaluate the performance of FIITED on four datasets, including three click-through rate prediction datasets and one classification dataset. Can you provide more details on the hyperparameters used in your experiments, such as the learning rate and batch size, and how they were selected?

The learning rates and batch sizes were grid searched with some guidance from previous research papers that used the same datasets. We selected hyperparameters that yielded accuracy comparable to existing research papers using the same datasets. The details are provided below:

| Dataset | Learning rate | Batch size |
|---|---|---|
| Criteo Kaggle | 1e-3 | 4096 |
| Criteo Terabyte | 1e-3 | 2048 |
| Avazu | 1e-3 | 4096 |
| MovieLens-20M | 1e-3 | 512 |

Reviewer comment (Novelty):
In Section 4.3, you compare the performance of FIITED with that of AdaEmbed, a state-of-the-art in-training embedding pruning method. Can you provide more details on how AdaEmbed works and how it differs from FIITED?

Author response:
**Q5**. In Section 4.3, you compare the performance of FIITED with that of AdaEmbed, a state-of-the-art in-training embedding pruning method. Can you provide more details on how AdaEmbed works and how it differs from FIITED?

AdaEmbed is a special case of FIITED where the number of chunks K is 1, i.e., every embedding is just 1 chunk. If this 1 chunk is pruned, then the embedding is pruned entirely. In AdaEmbed, an embedding vector is either totally unpruned or totally pruned.

FIITED is a generalization of AdaEmbed and allows an embedding to be pruned in a fine-grained way.

Reviewer comment (Evaluation):
In Section 5, you discuss the limitations of your approach, including the trade-offs between memory savings and model quality, and the limitations in terms of scalability and computational complexity. Can you provide more details on how these limitations might impact the applicability of your approach to large-scale datasets?

Author response:
**Q6**. In Section 5, you discuss the limitations of your approach, including the trade-offs between memory savings and model quality, and the limitations in terms of scalability and computational complexity. Can you provide more details on how these limitations might impact the applicability of your approach to large-scale datasets?

When it comes to scaling to large-scale datasets that requires a large number of embeddings, the model quality is limited by the amount of memory available if FIITED is used to fit the model into the on-device memory. If the amount of on-device memory is fixed, then there will be a point where the pruning rate needs to be so high that the model quality will suffer. Then, to ensure model quality, we will need to increase the amount of memory by, for example, increasing the number of devices used for training. A contribution of FIITED is that, given a memory budget, FIITED can figure out how to smartly prune embeddings to fit within the memory budget while achieving relatively good model quality compared to baselines in the field.

As for the performance overhead of FIITED in terms of training time, it is a tug of war between the extra computation introduced by FIITED (as discussed in Q1) and the amount of computation and data transfer saved by FIITED. Computation is saved by increasing the sparsity in data. If the model is too big to fit in on-device memory, expensive data transfer is usually needed to obtain embeddings from off-device memory, and FIITED can save such data transfer by drastically reducing the model size.

Please let us know if you have any questions or comments! Thank you.

Reviewer comment (Theory):
At the heart of FIITED is the utility-based approach to determine chunk significance. However, basing eviction decisions purely on utility scores might introduce biases. For instance, recent chunks might gain a temporary high utility, leading to potentially premature evictions of other valuable chunks.

Author response:
**Q1**. How does FIITED handle potential fluctuations in utility values?

FIITED takes history utility scores into consideration. As mentioned in Section 3.2, the utility value of a chunk in the i-th iteration is computed as follows:

u(i) = γu(i − 1) + a(i)g(i)

where u(i − 1) is the utility score in the previous iteration, a(i) is the number of accesses to the chunk in iteration i, g(i) is the L2 norm of the gradients computed for the chunk in iteration i, and γ ∈ (0, 1) is a decay parameter that controls the influence of the history utility value.

We want to make sure small fluctuations do not affect the big picture, and at the same time adapt to the dynamic changes in data characteristics. For example, if certain embeddings were "hot" in the past but no longer relevant in the present, it may be better to give the space away to another embedding that is currently accessed frequently.

As for the noisy gradients at the start of training, since it takes a while to fill up the embedding table before pruning needs to be triggered, the influence of noise is reduced.

Reviewer comment (Theory):
This approach does not consider the individual significance of dimensions within a chunk, leading to potential information loss.

Author response:
**Q2**. The authors may consider analyzing the dimensional dependency and significance first, as well as the adaptive chunking as a fixed chunk size may not always be optimal.

This is a really interesting point and an interesting direction for future work. I think information loss is almost inevitable as far as pruning is concerned. Structured model pruning, like FIITED, risks losing potentially important model parameters but often offers more efficient storage and computation, compared with unstructured pruning. However, since embedding dimensions are symmetric in the arithmetic sense (the input dimensions are symmetric in dot product and fully connected layers used in MLP - note that embeddings considered in this work only correspond to 1 feature at a time and all dimensions within an embedding are always accessed together), there can potentially be a way to use additional regularization terms to help form certain group structures in the dimensions so that information loss is reduced.

Reviewer comment (Theory):
While the chunk address manager maintains a free address stack, this design assumes that the most recently evicted space is optimal for the next allocation. This might not always be the case, especially when considering the locality of data and frequent access patterns.

Author response:
**Q3**. Instead of purely stack-based address management, the authors may consider using a hybrid approach that takes into account both recent evictions and access frequency.

Optimizing the reallocation address management policy can definitely improve the spatial locality, but temporal locality-wise, it does not affect the performance much because pruning does not happen frequently (e.g., every few hundreds of iterations) in FIITED. When pruning does happen, it is a short, concentrated time period when a group of less important embedding chunks are discarded and re-initialized to make space for another group of more important chunks. There is 1 swap per pruned chunk, and the free address stack typically ends up empty at the end of the pruning and reallocation process. It is possible that certain chunk locations will keep getting swapped, but every 2 swaps are many training iterations apart.

Reviewer comment (Theory):
The system heavily depends on the hash table to fetch and manage embeddings. This approach, while efficient in accessing chunks, might lead to hashing collisions even though the design ensures a low collision rate. Any collision, however rare, can introduce latency in access times or even potential overwrites.

Author response:
**Q4**. The authors may need to employ a more robust hashing mechanism or a secondary probing mechanism to handle collisions more effectively.

The collision issue is not unique to FIITED, as hashing is commonly used in modern Deep Learning Recommender Models (DLRM) to cut down the large number of embeddings. Smart hashing techniques have been proposed to deal with the collision problem in DLRM. Currently, FIITED simply uses a hash table with a fairly large number of rows to make sure collision rarely happens, but moving forward, it can definitely benefit from an improved hashing scheme.

Reviewer comment (Theory):
The methodology leans heavily on access frequency to decide on embedding significance. However, frequency doesn't always equate to importance. There could be rarely accessed but critically important embeddings, and the method might be prone to undervaluing them.

Author response:
**Q5**. Some dimensions might be accessed less frequently because the specific feature they capture is less common in the dataset. Yet, this doesn't necessarily mean that feature is unimportant.

It is true that less accessed features may be important to a small group of users, but the problem of fair recommendation is a quite broad subject that will be hard to solve in one paper. Existing embedding pruning papers generally aim to optimize for model accuracy while reducing the model size as much as possible, which means the broader user base will inevitably be favored due to them contributing to the majority of validation/test samples. However, FIITED can potentially help to facilitate fairness. By saving memory usage during training, FIITED can allow the addition of more features into the model to diversify the representation of user / item traits. Moreover, one can design better utility metrics to promote fairness among features.

Reviewer comment (Theory):
In the context of FIITED, access frequency influences the utility. When the system frequently accesses certain recent chunks, those chunks might temporarily have a higher utility score compared to older chunks—even if those older chunks are historically more important. Although FIITED utilizes the gradients of the chunk to mitigate such biases, the gradient might not always represent the "big picture" during the training process. For example, in the context of complex models, significant gradients might lead the optimization towards local minima rather than a more global optimal solution. Especially in the early stages of training, gradients can be noisy. This noise might mislead the system into thinking an embedding chunk is more or less important than it truly is. The authors should discuss more how it handles such potential fluctuations in utility values or if any smoothing mechanisms are in place.

Author response:
**Q1**. How does FIITED handle potential fluctuations in utility values?

FIITED takes history utility scores into consideration. As mentioned in Section 3.2, the utility value of a chunk in the i-th iteration is computed as follows:

u(i) = γu(i − 1) + a(i)g(i)

where u(i − 1) is the utility score in the previous iteration, a(i) is the number of accesses to the chunk in iteration i, g(i) is the L2 norm of the gradients computed for the chunk in iteration i, and γ ∈ (0, 1) is a decay parameter that controls the influence of the history utility value.

We want to make sure small fluctuations do not affect the big picture, and at the same time adapt to the dynamic changes in data characteristics. For example, if certain embeddings were "hot" in the past but no longer relevant in the present, it may be better to give the space away to another embedding that is currently accessed frequently.

As for the noisy gradients at the start of training, since it takes a while to fill up the embedding table before pruning needs to be triggered, the influence of noise is reduced.

Reviewer comment (Theory):
Dimensions in embeddings often have dependencies, where certain dimensions complement or refine the information in others. Chunking might break these dependencies. This can cause the model to lose subsets of interrelated features, compromising its understanding. Moreover, embedding dimensions might not have uniform importance. By dividing them into chunks, there's a risk that a chunk might have a mix of crucial and less vital dimensions. If the less important dimensions in a chunk increase its eviction likelihood, it might inadvertently lead to evicting more essential dimensions too. The authors may consider analyzing the dimensional dependency and significance first, as well as the adaptive chunking as a fixed chunk size may not always be optimal.

Author response:
**Q2**. The authors may consider analyzing the dimensional dependency and significance first, as well as the adaptive chunking as a fixed chunk size may not always be optimal.

This is a really interesting point and an interesting direction for future work. I think information loss is almost inevitable as far as pruning is concerned. Structured model pruning, like FIITED, risks losing potentially important model parameters but often offers more efficient storage and computation, compared with unstructured pruning. However, since embedding dimensions are symmetric in the arithmetic sense (the input dimensions are symmetric in dot product and fully connected layers used in MLP - note that embeddings considered in this work only correspond to 1 feature at a time and all dimensions within an embedding are always accessed together), there can potentially be a way to use additional regularization terms to help form certain group structures in the dimensions so that information loss is reduced.

Reviewer comment (Theory):
The address manager simply re-allocates the most recently evicted space (due to the stack-based approach), and it might re-allocate space that has high temporal locality. Thus, it could be evicting data that will be frequently accessed shortly after. This isn't efficient, as it could result in a high rate of data swapping in and out. Hence, simply using a stack to manage the evicted spaces could disrupt both forms of locality. If the system doesn't take into account the access patterns and just works on a "most recently freed" basis, it could lead to inefficiencies in memory access and potential performance bottlenecks, especially in scenarios where embeddings have strong temporal or spatial access patterns. Instead of purely stack-based address management, the authors may consider using a hybrid approach that takes into account both recent evictions and access frequency.

Author response:
**Q3**. Instead of purely stack-based address management, the authors may consider using a hybrid approach that takes into account both recent evictions and access frequency.

Optimizing the reallocation address management policy can definitely improve the spatial locality, but temporal locality-wise, it does not affect the performance much because pruning does not happen frequently (e.g., every few hundreds of iterations) in FIITED. When pruning does happen, it is a short, concentrated time period when a group of less important embedding chunks are discarded and re-initialized to make space for another group of more important chunks. There is 1 swap per pruned chunk, and the free address stack typically ends up empty at the end of the pruning and reallocation process. It is possible that certain chunk locations will keep getting swapped, but every 2 swaps are many training iterations apart.

Reviewer comment (Evaluation):
Unlike whole embeddings that represent specific entities (like items), dimensions within embeddings capture specific features or facets of the represented entities. Some dimensions might be accessed less frequently because the specific feature they capture is less common in the dataset. Yet, this doesn't necessarily mean that feature is unimportant. For example, in e-commerce recommendation systems, there's often a "long tail" of products that aren't as popular as the mainstream ones but cater to specific needs or tastes. The dimensions representing these products might not be accessed as often, but they can be key to offering diverse and personalized recommendations to users looking for something unique. In addition, some users might have niche interests that aren't shared by the broader user base. The dimensions representing these niche interests might be accessed less frequently due to the smaller number of users with these interests. However, for those specific users, these dimensions are essential for generating accurate and satisfying recommendations.

Author response:
**Q5**. Some dimensions might be accessed less frequently because the specific feature they capture is less common in the dataset. Yet, this doesn't necessarily mean that feature is unimportant.

It is true that less accessed features may be important to a small group of users, but the problem of fair recommendation is a quite broad subject that will be hard to solve in one paper. Existing embedding pruning papers generally aim to optimize for model accuracy while reducing the model size as much as possible, which means the broader user base will inevitably be favored due to them contributing to the majority of validation/test samples. However, FIITED can potentially help to facilitate fairness. By saving memory usage during training, FIITED can allow the addition of more features into the model to diversify the representation of user / item traits. Moreover, one can design better utility metrics to promote fairness among features.