Transformer-based cross-node machine learning systems for wireless communication
Various aspects of the present disclosure generally relate to wireless communication. In some aspects, a user equipment (UE) may receive a transformer configuration that includes a transmitter neural network configured to be used to generate at least one latent vector corresponding to one or more computation tasks of a plurality of computation tasks associated with a transformer-based cross-node machine learning system. The UE may transmit the at least one latent vector based at least in part on instantiating the transmitter neural network. Numerous other aspects are described.
1. A user equipment (UE) for wireless communication, comprising:
one or more memories; and
one or more processors, coupled to the one or more memories, configured to:
receive a transformer configuration that includes a transmitter neural network configured to be used to generate at least one latent vector corresponding to one or more computation tasks of a plurality of computation tasks associated with a transformer-based cross-node machine learning system, wherein the transformer-based cross-node machine learning system comprises the transmitter neural network instantiated by the UE; and
transmit the at least one latent vector based at least in part on instantiating the transmitter neural network.
2. The UE of claim 1 , wherein the transformer-based cross-node machine learning system further comprises a plurality of receiver neural networks instantiated by a base station, wherein each receiver neural network of the plurality of receiver neural networks corresponds to a computation task of the plurality of computation tasks.
3. The UE of claim 1 , wherein the transmitter neural network comprises a transformer backbone instance corresponding to the transformer-based cross-node machine learning system.
4. The UE of claim 3 , wherein the transformer backbone instance comprises:
a linear projection component that takes, as input, a set of input tokens and generates a set of linear token embeddings corresponding to the set of input tokens, respectively;
a transmitter positional encoding component that takes, as input, the set of linear token embeddings and a set of task embedding vectors, wherein each task embedding vector of the set of task embedding vectors corresponds to one of the one or more computation tasks, and wherein the transmitter positional encoding component generates a set of token embedding vectors corresponding to the set of linear token embeddings and a set of position-encoded task embedding vectors corresponding to the set of task embedding vectors; and
a transmitter transformer encoder that takes, as input, the set of token embedding vectors and the set of position-encoded task embedding vectors, wherein the transmitter transformer encoder generates a set of transformed token embedding vectors corresponding to the set of token embedding vectors and a set of transformed task embedding vectors corresponding to the set of position-encoded task embedding vectors.
5. The UE of claim 1 , wherein the transmitter neural network comprises one or more layers of at least one head module instance corresponding to the one or more computation tasks.
6. The UE of claim 1 , wherein the latent vector comprises an output of one or more layers of a head module instance corresponding to a computation task of the one or more computation tasks, wherein the head module instance takes, as input, a transformed task embedding vector generated by a transmitter transformer encoder.
7. The UE of claim 1 , wherein the one or more processors, to receive the transformer configuration, are configured to receive an upper-layer communication including the transformer configuration.
8. The UE of claim 7 , wherein the upper-layer communication comprises a radio resource control message.
9. The UE of claim 1 , wherein the transformer configuration indicates at least one of:
a set of transmitter transformer encoder parameters,
a position embedding matrix,
a linear projection matrix,
a set of task embedding vectors,
an indication of an ordering of the set of task embedding vectors and a set of linear token embeddings, or
a set of head component parameters.
10. The UE of claim 1 , wherein the transformer configuration indicates a reporting configuration, and wherein the one or more processors, to transmit the at least one latent vector, are configured to transmit the at least one latent vector based at least in part on the reporting configuration.
11. The UE of claim 10 , wherein the reporting configuration comprises at least one of:
an indication of a reporting period,
an indication of a reporting offset time, or
an indication of a task-based reporting scheme.
12. The UE of claim 1 , wherein the one or more processors are further configured to receive a downlink control information (DCI) transmission that includes an indication to report the at least one latent vector, and wherein the one or more processors, to transmit the at least one latent vector, are configured to transmit the at least one latent vector based at least in part on receiving the DCI transmission.
13. A network entity for wireless communication, comprising:
one or more memories; and
one or more processors, coupled to the one or more memories, configured to:
transmit, to a user equipment (UE), a transformer configuration that includes a transmitter neural network configured to be used to generate at least one latent vector corresponding to one or more computation tasks of a plurality of computation tasks associated with a transformer-based cross-node machine learning system, wherein the transformer-based cross-node machine learning system comprises the transmitter neural network instantiated by the UE; and
receive, from the UE and based at least in part on instantiation of the transmitter neural network by the UE, the at least one latent vector.
14. The network entity of claim 13 , wherein the one or more processors are further configured to:
transmit, to at least one additional UE, at least one additional transformer configuration that includes at least one additional transmitter neural network configured to be used to generate at least one additional latent vector corresponding to one or more additional computation tasks of the plurality of computation tasks associated with the transformer-based cross-node machine learning system; and
receive, from the at least one additional UE, the at least one additional latent vector.
15. The network entity of claim 13 , wherein the transformer-based cross-node machine learning system further comprises a receiver neural network instantiated by the network entity.
16. The network entity of claim 13 , wherein the transformer-based cross-node machine learning system further comprises a plurality of receiver neural networks instantiated by the network entity, wherein each receiver neural network of the plurality of receiver neural networks corresponds to a computation task of the plurality of computation tasks.
17. The network entity of claim 16 , wherein the transmitter neural network comprises one or more layers of a transmitter instance of a head module corresponding to a computation task of the one or more computation tasks, and wherein a receiver neural network of the plurality of receiver neural networks comprises one or more additional layers of a receiver instance of the head module.
18. The network entity of claim 13 , wherein the transformer-based cross-node machine learning system further comprises at least one additional transmitter neural network instantiated by at least one additional UE.
19. The network entity of claim 13 , wherein the one or more processors are further configured to provide the latent vector as input to a receiver neural network corresponding to a computation task of the plurality of computation tasks.
20. The network entity of claim 19 , wherein the receiver neural network comprises:
a mapping head module instance corresponding to the computation task, wherein the mapping head module instance maps the latent vector to a mapped embedding vector;
a receiver positional encoding component that takes, as input, the mapped embedding vector and a set of zero vectors, and wherein the receiver positional encoding component generates a set of token queries corresponding to the set of zero vectors and a position-encoded mapped embedding vector corresponding to the mapped embedding vector;
a receiver transformer encoder that takes, as input, the set of token queries and the position-encoded mapped embedding vector, wherein the receiver transformer encoder generates a set of unpacked token vectors corresponding to the set of token queries; and
a set of estimation head modules that take, as input, the set of unpacked token vectors, wherein each estimation head module of the set of estimation head modules generates an estimated input token.
21. The network entity of claim 13 , wherein the at least one latent vector comprises an output of one or more layers of a head module instance corresponding to a computation task of the one or more computation tasks.
22. The network entity of claim 13 , wherein the one or more processors, to transmit the transformer configuration, are configured to transmit an upper-layer communication including the transformer configuration.
23. The network entity of claim 22 , wherein the upper-layer communication comprises a radio resource control message.
24. The network entity of claim 13 , wherein the transformer configuration indicates at least one of:
a set of transmitter transformer encoder parameters,
a position embedding matrix,
a linear projection matrix,
a set of task embedding vectors,
an indication of an ordering of the set of task embedding vectors and a set of linear token embeddings, or
a set of head component parameters.
25. The network entity of claim 13 , wherein the transformer configuration indicates a reporting configuration, and wherein the one or more processors, to receive the at least one latent vector, are configured to receive the at least one latent vector based at least in part on the reporting configuration.
26. The network entity of claim 13 , wherein the one or more processors are further configured to transmit a downlink control information (DCI) transmission that includes an indication to report the at least one latent vector, and wherein the one or more processors, to receive the at least one latent vector, are configured to receive the at least one latent vector based at least in part on transmitting the DCI transmission.
27. A method of wireless communication performed by a user equipment (UE), comprising:
receiving a transformer configuration that includes a transmitter neural network configured to be used to generate at least one latent vector corresponding to one or more computation tasks of a plurality of computation tasks associated with a transformer-based cross-node machine learning system, wherein the transformer-based cross-node machine learning system comprises the transmitter neural network instantiated by the UE; and
transmitting the at least one latent vector based at least in part on instantiating the transmitter neural network.
28. The method of claim 27 , wherein the transformer-based cross-node machine learning system further comprises a plurality of receiver neural networks instantiated by a network entity, each receiver neural network of the plurality of receiver neural networks corresponding to a computation task of the plurality of computation tasks, wherein receiving the transformer configuration comprises receiving an upper-layer communication including the transformer configuration, and wherein the transformer configuration indicates at least one of:
a set of transmitter transformer encoder parameters,
a position embedding matrix,
a linear projection matrix,
a set of task embedding vectors,
an indication of an ordering of the set of task embedding vectors and a set of linear token embeddings, or
a set of head component parameters.
29. A method of wireless communication performed by a network entity, comprising:
transmitting, to a user equipment (UE), a transformer configuration that includes a transmitter neural network configured to be used to generate at least one latent vector corresponding to one or more computation tasks of a plurality of computation tasks associated with a transformer-based cross-node machine learning system, wherein the transformer-based cross-node machine learning system comprises the transmitter neural network instantiated by the UE; and
receiving, from the UE and based at least in part on instantiation of the transmitter neural network by the UE, the at least one latent vector.
30. The method of claim 29 , further comprising:
transmitting, to at least one additional UE, at least one additional transformer configuration that includes at least one additional transmitter neural network configured to be used to generate at least one additional latent vector corresponding to one or more additional computation tasks of the plurality of computation tasks associated with the transformer-based cross-node machine learning system; and
receiving, from the at least one additional UE, the at least one additional latent vector.