├── .gitmodules ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── compute_renyi_div.py ├── install.py ├── mean_estimation_multi.py ├── mean_estimation_single.py ├── mechanisms.py ├── mechanisms_pytorch.py ├── optimize_mvu.py ├── patches └── private_prediction.patch ├── plot_dme_l1.py ├── plot_dme_l2.py ├── requirements.txt ├── train.py └── utils.py /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "private_prediction"] 2 | path = private_prediction 3 | url = https://github.com/facebookresearch/private_prediction 4 | [submodule "Handcrafted-DP"] 5 | path = Handcrafted-DP 6 | url = https://github.com/ftramer/Handcrafted-DP 7 | [submodule "fastwht"] 8 | path = fastwht 9 | url = git@github.com:vegarant/fastwht.git 10 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to make participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, sex characteristics, gender identity and expression, 9 | level of experience, education, socio-economic status, nationality, personal 10 | appearance, race, religion, or sexual identity and orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | * Using welcoming and inclusive language 18 | * Being respectful of differing viewpoints and experiences 19 | * Gracefully accepting constructive criticism 20 | * Focusing on what is best for the community 21 | * Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | * The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | * Trolling, insulting/derogatory comments, and personal or political attacks 28 | * Public or private harassment 29 | * Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | * Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or 41 | reject comments, commits, code, wiki edits, issues, and other contributions 42 | that are not aligned to this Code of Conduct, or to ban temporarily or 43 | permanently any contributor for other behaviors that they deem inappropriate, 44 | threatening, offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies within all project spaces, and it also applies when 49 | an individual is representing the project or its community in public spaces. 50 | Examples of representing a project or community include using an official 51 | project e-mail address, posting via an official social media account, or acting 52 | as an appointed representative at an online or offline event. Representation of 53 | a project may be further defined and clarified by project maintainers. 54 | 55 | ## Enforcement 56 | 57 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 58 | reported by contacting the project team at . All 59 | complaints will be reviewed and investigated and will result in a response that 60 | is deemed necessary and appropriate to the circumstances. The project team is 61 | obligated to maintain confidentiality with regard to the reporter of an incident. 62 | Further details of specific enforcement policies may be posted separately. 63 | 64 | Project maintainers who do not follow or enforce the Code of Conduct in good 65 | faith may face temporary or permanent repercussions as determined by other 66 | members of the project's leadership. 67 | 68 | ## Attribution 69 | 70 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, 71 | available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html 72 | 73 | [homepage]: https://www.contributor-covenant.org 74 | 75 | For answers to common questions about this code of conduct, see 76 | https://www.contributor-covenant.org/faq 77 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to Privacy-Aware Data Compression 2 | We want to make contributing to this project as easy and transparent as 3 | possible. 4 | 5 | ## Our Development Process 6 | ... (in particular how this is synced with internal changes to the project) 7 | 8 | ## Pull Requests 9 | We actively welcome your pull requests. 10 | 11 | 1. Fork the repo and create your branch from `main`. 12 | 2. If you've added code that should be tested, add tests. 13 | 3. If you've changed APIs, update the documentation. 14 | 4. Ensure the test suite passes. 15 | 5. Make sure your code lints. 16 | 6. If you haven't already, complete the Contributor License Agreement ("CLA"). 17 | 18 | ## Contributor License Agreement ("CLA") 19 | In order to accept your pull request, we need you to submit a CLA. You only need 20 | to do this once to work on any of Meta's open source projects. 21 | 22 | Complete your CLA here: 23 | 24 | ## Issues 25 | We use GitHub issues to track public bugs. Please ensure your description is 26 | clear and has sufficient instructions to be able to reproduce the issue. 27 | 28 | Meta has a [bounty program](https://www.facebook.com/whitehat/) for the safe 29 | disclosure of security bugs. In those cases, please go through the process 30 | outlined on that page and do not file a public issue. 31 | 32 | ## Coding Style 33 | * 2 spaces for indentation rather than tabs 34 | * 80 character line length 35 | * ... 36 | 37 | ## License 38 | By contributing to Privacy-Aware Data Compression, you agree that your contributions will be licensed 39 | under the LICENSE file in the root directory of this source tree. 40 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Attribution-NonCommercial 4.0 International 3 | 4 | ======================================================================= 5 | 6 | Creative Commons Corporation ("Creative Commons") is not a law firm and 7 | does not provide legal services or legal advice. Distribution of 8 | Creative Commons public licenses does not create a lawyer-client or 9 | other relationship. Creative Commons makes its licenses and related 10 | information available on an "as-is" basis. Creative Commons gives no 11 | warranties regarding its licenses, any material licensed under their 12 | terms and conditions, or any related information. Creative Commons 13 | disclaims all liability for damages resulting from their use to the 14 | fullest extent possible. 15 | 16 | Using Creative Commons Public Licenses 17 | 18 | Creative Commons public licenses provide a standard set of terms and 19 | conditions that creators and other rights holders may use to share 20 | original works of authorship and other material subject to copyright 21 | and certain other rights specified in the public license below. The 22 | following considerations are for informational purposes only, are not 23 | exhaustive, and do not form part of our licenses. 24 | 25 | Considerations for licensors: Our public licenses are 26 | intended for use by those authorized to give the public 27 | permission to use material in ways otherwise restricted by 28 | copyright and certain other rights. Our licenses are 29 | irrevocable. Licensors should read and understand the terms 30 | and conditions of the license they choose before applying it. 31 | Licensors should also secure all rights necessary before 32 | applying our licenses so that the public can reuse the 33 | material as expected. Licensors should clearly mark any 34 | material not subject to the license. This includes other CC- 35 | licensed material, or material used under an exception or 36 | limitation to copyright. More considerations for licensors: 37 | wiki.creativecommons.org/Considerations_for_licensors 38 | 39 | Considerations for the public: By using one of our public 40 | licenses, a licensor grants the public permission to use the 41 | licensed material under specified terms and conditions. If 42 | the licensor's permission is not necessary for any reason--for 43 | example, because of any applicable exception or limitation to 44 | copyright--then that use is not regulated by the license. Our 45 | licenses grant only permissions under copyright and certain 46 | other rights that a licensor has authority to grant. Use of 47 | the licensed material may still be restricted for other 48 | reasons, including because others have copyright or other 49 | rights in the material. A licensor may make special requests, 50 | such as asking that all changes be marked or described. 51 | Although not required by our licenses, you are encouraged to 52 | respect those requests where reasonable. More_considerations 53 | for the public: 54 | wiki.creativecommons.org/Considerations_for_licensees 55 | 56 | ======================================================================= 57 | 58 | Creative Commons Attribution-NonCommercial 4.0 International Public 59 | License 60 | 61 | By exercising the Licensed Rights (defined below), You accept and agree 62 | to be bound by the terms and conditions of this Creative Commons 63 | Attribution-NonCommercial 4.0 International Public License ("Public 64 | License"). To the extent this Public License may be interpreted as a 65 | contract, You are granted the Licensed Rights in consideration of Your 66 | acceptance of these terms and conditions, and the Licensor grants You 67 | such rights in consideration of benefits the Licensor receives from 68 | making the Licensed Material available under these terms and 69 | conditions. 70 | 71 | Section 1 -- Definitions. 72 | 73 | a. Adapted Material means material subject to Copyright and Similar 74 | Rights that is derived from or based upon the Licensed Material 75 | and in which the Licensed Material is translated, altered, 76 | arranged, transformed, or otherwise modified in a manner requiring 77 | permission under the Copyright and Similar Rights held by the 78 | Licensor. For purposes of this Public License, where the Licensed 79 | Material is a musical work, performance, or sound recording, 80 | Adapted Material is always produced where the Licensed Material is 81 | synched in timed relation with a moving image. 82 | 83 | b. Adapter's License means the license You apply to Your Copyright 84 | and Similar Rights in Your contributions to Adapted Material in 85 | accordance with the terms and conditions of this Public License. 86 | 87 | c. Copyright and Similar Rights means copyright and/or similar rights 88 | closely related to copyright including, without limitation, 89 | performance, broadcast, sound recording, and Sui Generis Database 90 | Rights, without regard to how the rights are labeled or 91 | categorized. For purposes of this Public License, the rights 92 | specified in Section 2(b)(1)-(2) are not Copyright and Similar 93 | Rights. 94 | d. Effective Technological Measures means those measures that, in the 95 | absence of proper authority, may not be circumvented under laws 96 | fulfilling obligations under Article 11 of the WIPO Copyright 97 | Treaty adopted on December 20, 1996, and/or similar international 98 | agreements. 99 | 100 | e. Exceptions and Limitations means fair use, fair dealing, and/or 101 | any other exception or limitation to Copyright and Similar Rights 102 | that applies to Your use of the Licensed Material. 103 | 104 | f. Licensed Material means the artistic or literary work, database, 105 | or other material to which the Licensor applied this Public 106 | License. 107 | 108 | g. Licensed Rights means the rights granted to You subject to the 109 | terms and conditions of this Public License, which are limited to 110 | all Copyright and Similar Rights that apply to Your use of the 111 | Licensed Material and that the Licensor has authority to license. 112 | 113 | h. Licensor means the individual(s) or entity(ies) granting rights 114 | under this Public License. 115 | 116 | i. NonCommercial means not primarily intended for or directed towards 117 | commercial advantage or monetary compensation. For purposes of 118 | this Public License, the exchange of the Licensed Material for 119 | other material subject to Copyright and Similar Rights by digital 120 | file-sharing or similar means is NonCommercial provided there is 121 | no payment of monetary compensation in connection with the 122 | exchange. 123 | 124 | j. Share means to provide material to the public by any means or 125 | process that requires permission under the Licensed Rights, such 126 | as reproduction, public display, public performance, distribution, 127 | dissemination, communication, or importation, and to make material 128 | available to the public including in ways that members of the 129 | public may access the material from a place and at a time 130 | individually chosen by them. 131 | 132 | k. Sui Generis Database Rights means rights other than copyright 133 | resulting from Directive 96/9/EC of the European Parliament and of 134 | the Council of 11 March 1996 on the legal protection of databases, 135 | as amended and/or succeeded, as well as other essentially 136 | equivalent rights anywhere in the world. 137 | 138 | l. You means the individual or entity exercising the Licensed Rights 139 | under this Public License. Your has a corresponding meaning. 140 | 141 | Section 2 -- Scope. 142 | 143 | a. License grant. 144 | 145 | 1. Subject to the terms and conditions of this Public License, 146 | the Licensor hereby grants You a worldwide, royalty-free, 147 | non-sublicensable, non-exclusive, irrevocable license to 148 | exercise the Licensed Rights in the Licensed Material to: 149 | 150 | a. reproduce and Share the Licensed Material, in whole or 151 | in part, for NonCommercial purposes only; and 152 | 153 | b. produce, reproduce, and Share Adapted Material for 154 | NonCommercial purposes only. 155 | 156 | 2. Exceptions and Limitations. For the avoidance of doubt, where 157 | Exceptions and Limitations apply to Your use, this Public 158 | License does not apply, and You do not need to comply with 159 | its terms and conditions. 160 | 161 | 3. Term. The term of this Public License is specified in Section 162 | 6(a). 163 | 164 | 4. Media and formats; technical modifications allowed. The 165 | Licensor authorizes You to exercise the Licensed Rights in 166 | all media and formats whether now known or hereafter created, 167 | and to make technical modifications necessary to do so. The 168 | Licensor waives and/or agrees not to assert any right or 169 | authority to forbid You from making technical modifications 170 | necessary to exercise the Licensed Rights, including 171 | technical modifications necessary to circumvent Effective 172 | Technological Measures. For purposes of this Public License, 173 | simply making modifications authorized by this Section 2(a) 174 | (4) never produces Adapted Material. 175 | 176 | 5. Downstream recipients. 177 | 178 | a. Offer from the Licensor -- Licensed Material. Every 179 | recipient of the Licensed Material automatically 180 | receives an offer from the Licensor to exercise the 181 | Licensed Rights under the terms and conditions of this 182 | Public License. 183 | 184 | b. No downstream restrictions. You may not offer or impose 185 | any additional or different terms or conditions on, or 186 | apply any Effective Technological Measures to, the 187 | Licensed Material if doing so restricts exercise of the 188 | Licensed Rights by any recipient of the Licensed 189 | Material. 190 | 191 | 6. No endorsement. Nothing in this Public License constitutes or 192 | may be construed as permission to assert or imply that You 193 | are, or that Your use of the Licensed Material is, connected 194 | with, or sponsored, endorsed, or granted official status by, 195 | the Licensor or others designated to receive attribution as 196 | provided in Section 3(a)(1)(A)(i). 197 | 198 | b. Other rights. 199 | 200 | 1. Moral rights, such as the right of integrity, are not 201 | licensed under this Public License, nor are publicity, 202 | privacy, and/or other similar personality rights; however, to 203 | the extent possible, the Licensor waives and/or agrees not to 204 | assert any such rights held by the Licensor to the limited 205 | extent necessary to allow You to exercise the Licensed 206 | Rights, but not otherwise. 207 | 208 | 2. Patent and trademark rights are not licensed under this 209 | Public License. 210 | 211 | 3. To the extent possible, the Licensor waives any right to 212 | collect royalties from You for the exercise of the Licensed 213 | Rights, whether directly or through a collecting society 214 | under any voluntary or waivable statutory or compulsory 215 | licensing scheme. In all other cases the Licensor expressly 216 | reserves any right to collect such royalties, including when 217 | the Licensed Material is used other than for NonCommercial 218 | purposes. 219 | 220 | Section 3 -- License Conditions. 221 | 222 | Your exercise of the Licensed Rights is expressly made subject to the 223 | following conditions. 224 | 225 | a. Attribution. 226 | 227 | 1. If You Share the Licensed Material (including in modified 228 | form), You must: 229 | 230 | a. retain the following if it is supplied by the Licensor 231 | with the Licensed Material: 232 | 233 | i. identification of the creator(s) of the Licensed 234 | Material and any others designated to receive 235 | attribution, in any reasonable manner requested by 236 | the Licensor (including by pseudonym if 237 | designated); 238 | 239 | ii. a copyright notice; 240 | 241 | iii. a notice that refers to this Public License; 242 | 243 | iv. a notice that refers to the disclaimer of 244 | warranties; 245 | 246 | v. a URI or hyperlink to the Licensed Material to the 247 | extent reasonably practicable; 248 | 249 | b. indicate if You modified the Licensed Material and 250 | retain an indication of any previous modifications; and 251 | 252 | c. indicate the Licensed Material is licensed under this 253 | Public License, and include the text of, or the URI or 254 | hyperlink to, this Public License. 255 | 256 | 2. You may satisfy the conditions in Section 3(a)(1) in any 257 | reasonable manner based on the medium, means, and context in 258 | which You Share the Licensed Material. For example, it may be 259 | reasonable to satisfy the conditions by providing a URI or 260 | hyperlink to a resource that includes the required 261 | information. 262 | 263 | 3. If requested by the Licensor, You must remove any of the 264 | information required by Section 3(a)(1)(A) to the extent 265 | reasonably practicable. 266 | 267 | 4. If You Share Adapted Material You produce, the Adapter's 268 | License You apply must not prevent recipients of the Adapted 269 | Material from complying with this Public License. 270 | 271 | Section 4 -- Sui Generis Database Rights. 272 | 273 | Where the Licensed Rights include Sui Generis Database Rights that 274 | apply to Your use of the Licensed Material: 275 | 276 | a. for the avoidance of doubt, Section 2(a)(1) grants You the right 277 | to extract, reuse, reproduce, and Share all or a substantial 278 | portion of the contents of the database for NonCommercial purposes 279 | only; 280 | 281 | b. if You include all or a substantial portion of the database 282 | contents in a database in which You have Sui Generis Database 283 | Rights, then the database in which You have Sui Generis Database 284 | Rights (but not its individual contents) is Adapted Material; and 285 | 286 | c. You must comply with the conditions in Section 3(a) if You Share 287 | all or a substantial portion of the contents of the database. 288 | 289 | For the avoidance of doubt, this Section 4 supplements and does not 290 | replace Your obligations under this Public License where the Licensed 291 | Rights include other Copyright and Similar Rights. 292 | 293 | Section 5 -- Disclaimer of Warranties and Limitation of Liability. 294 | 295 | a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE 296 | EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS 297 | AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF 298 | ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS, 299 | IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION, 300 | WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR 301 | PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS, 302 | ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT 303 | KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT 304 | ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU. 305 | 306 | b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE 307 | TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION, 308 | NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT, 309 | INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES, 310 | COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR 311 | USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN 312 | ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR 313 | DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR 314 | IN PART, THIS LIMITATION MAY NOT APPLY TO YOU. 315 | 316 | c. The disclaimer of warranties and limitation of liability provided 317 | above shall be interpreted in a manner that, to the extent 318 | possible, most closely approximates an absolute disclaimer and 319 | waiver of all liability. 320 | 321 | Section 6 -- Term and Termination. 322 | 323 | a. This Public License applies for the term of the Copyright and 324 | Similar Rights licensed here. However, if You fail to comply with 325 | this Public License, then Your rights under this Public License 326 | terminate automatically. 327 | 328 | b. Where Your right to use the Licensed Material has terminated under 329 | Section 6(a), it reinstates: 330 | 331 | 1. automatically as of the date the violation is cured, provided 332 | it is cured within 30 days of Your discovery of the 333 | violation; or 334 | 335 | 2. upon express reinstatement by the Licensor. 336 | 337 | For the avoidance of doubt, this Section 6(b) does not affect any 338 | right the Licensor may have to seek remedies for Your violations 339 | of this Public License. 340 | 341 | c. For the avoidance of doubt, the Licensor may also offer the 342 | Licensed Material under separate terms or conditions or stop 343 | distributing the Licensed Material at any time; however, doing so 344 | will not terminate this Public License. 345 | 346 | d. Sections 1, 5, 6, 7, and 8 survive termination of this Public 347 | License. 348 | 349 | Section 7 -- Other Terms and Conditions. 350 | 351 | a. The Licensor shall not be bound by any additional or different 352 | terms or conditions communicated by You unless expressly agreed. 353 | 354 | b. Any arrangements, understandings, or agreements regarding the 355 | Licensed Material not stated herein are separate from and 356 | independent of the terms and conditions of this Public License. 357 | 358 | Section 8 -- Interpretation. 359 | 360 | a. For the avoidance of doubt, this Public License does not, and 361 | shall not be interpreted to, reduce, limit, restrict, or impose 362 | conditions on any use of the Licensed Material that could lawfully 363 | be made without permission under this Public License. 364 | 365 | b. To the extent possible, if any provision of this Public License is 366 | deemed unenforceable, it shall be automatically reformed to the 367 | minimum extent necessary to make it enforceable. If the provision 368 | cannot be reformed, it shall be severed from this Public License 369 | without affecting the enforceability of the remaining terms and 370 | conditions. 371 | 372 | c. No term or condition of this Public License will be waived and no 373 | failure to comply consented to unless expressly agreed to by the 374 | Licensor. 375 | 376 | d. Nothing in this Public License constitutes or may be interpreted 377 | as a limitation upon, or waiver of, any privileges and immunities 378 | that apply to the Licensor or You, including from the legal 379 | processes of any jurisdiction or authority. 380 | 381 | ======================================================================= 382 | 383 | Creative Commons is not a party to its public 384 | licenses. Notwithstanding, Creative Commons may elect to apply one of 385 | its public licenses to material it publishes and in those instances 386 | will be considered the “Licensor.” The text of the Creative Commons 387 | public licenses is dedicated to the public domain under the CC0 Public 388 | Domain Dedication. Except for the limited purpose of indicating that 389 | material is shared under a Creative Commons public license or as 390 | otherwise permitted by the Creative Commons policies published at 391 | creativecommons.org/policies, Creative Commons does not authorize the 392 | use of the trademark "Creative Commons" or any other trademark or logo 393 | of Creative Commons without its prior written consent including, 394 | without limitation, in connection with any unauthorized modifications 395 | to any of its public licenses or any other arrangements, 396 | understandings, or agreements concerning use of licensed material. For 397 | the avoidance of doubt, this paragraph does not form part of the 398 | public licenses. 399 | 400 | Creative Commons may be contacted at creativecommons.org. 401 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Privacy-Aware Compression for Federated Data Analysis 2 | 3 | This repository contains code for reproducing results in the papers: 4 | - Kamalika Chaudhuri*, Chuan Guo*, Mike Rabbat. **[Privacy-Aware Compression for Federated Data Analysis](https://arxiv.org/abs/2203.08134)**. 5 | - Chuan Guo, Kamalika Chaudhuri, Pierre Stock, Mike Rabbat. **[Privacy-Aware Compression for Federated Learning Through Numerical Mechanism Design](https://arxiv.org/abs/2211.03942)**. 6 | 7 | ## Setup 8 | 9 | Dependencies: [numpy](https://numpy.org/), [scipy](https://scipy.org/), [cvxpy](https://www.cvxpy.org/), [pytorch](https://pytorch.org/), [opacus](https://github.com/pytorch/opacus), [kymatio](https://github.com/kymatio/kymatio), [Handcrafted-DP](https://github.com/ftramer/Handcrafted-DP), [private_prediction](https://github.com/facebookresearch/private_prediction), [fastwht](https://github.com/vegarant/fastwht). 10 | 11 | After cloning repo and installing dependencies (see `requirements.txt`), download submodules and run the install script to apply some patches. 12 | ``` 13 | git submodule update --init 14 | python install.py 15 | cd fastwht/python 16 | ./setup.sh 17 | ``` 18 | 19 | ## Experiments 20 | 21 | ### Scalar Distributed Mean Estimation (CPU only) 22 | 23 | ``` 24 | for epsilon in 1 3 5; do 25 | python optimize_mvu.py --input_bits 3 --budget 3 --epsilon $epsilon --dp_constraint strict --method tr 26 | python mean_estimation_single.py --epsilon $epsilon 27 | done 28 | ``` 29 | 30 | ### Vector Distributed Mean Estimation (CPU only) 31 | 32 | For L1-sensitivity setting, first optimize the MVU mechanisms: 33 | ``` 34 | for epsilon in 1 2 3 4 5 6 7 8 9 10; do 35 | python optimize_mvu.py --input_bits 9 --budget 3 --epsilon $epsilon --dp_constraint metric-l1 --method penalized 36 | done 37 | ``` 38 | Then run the DME experiment and plot result: 39 | ``` 40 | for epsilon in 0.5 1 1.5 2 2.5 3 3.5 4 4.5 5; do 41 | python mean_estimation_multi.py --norm_type l1 --epsilon $epsilon --skellam_budget 16 --skellam_s 100 --mvu_input_bits 9 --mvu_budget 3 42 | done 43 | python plot_dme_l1.py 44 | ``` 45 | 46 | For L2-sensitivity setting, first optimize the MVU mechanisms and compute Renyi divergence curve for both the pure and approximate DP variants: 47 | ``` 48 | for epsilon in 2 4 6 8 10 12 14 16 18 20; do 49 | python optimize_mvu.py --input_bits 5 --budget 3 --epsilon $epsilon --dp_constraint metric-l2 --method penalized 50 | done 51 | for epsilon in 0.25 0.5 0.75 1 1.25 1.5 1.75 2 2.25 2.5; do 52 | python optimize_mvu.py --input_bits 5 --budget 3 --epsilon $epsilon --dp_constraint metric-l1 --method penalized 53 | python compute_renyi_div.py --input_bits 5 --budget 3 --epsilon $epsilon --dp_constraint metric-l1 54 | done 55 | ``` 56 | Then run the DME experiment and plot result: 57 | ``` 58 | for epsilon in 0.5 1 1.5 2 2.5 3 3.5 4 4.5 5; do 59 | python mean_estimation_multi.py --norm_type l2 --epsilon $epsilon --skellam_budget 16 --skellam_s 15 --mvu_input_bits 5 --mvu_budget 3 60 | done 61 | python plot_dme_l2.py 62 | ``` 63 | 64 | **Note**: These experiments will take a few hours to run. 65 | 66 | ### DP-SGD Training (requires GPU) 67 | 68 | 69 | To run the DP-SGD training experiment, first optimize the MVU mechanism: 70 | ``` 71 | python optimize_mvu.py --input_bits 9 --budget 1 --epsilon --dp_constraint metric-l1 --method penalized 72 | python optimize_mvu.py --input_bits 1 --budget 1 --epsilon --dp_constraint metric-l1 --method penalized 73 | ``` 74 | Then run DP-SGD training with Gaussian mechanism, signSGD, Skellam, MVU and I-MVU: 75 | ``` 76 | python train.py --save-model --dataset mnist --model --mechanism gaussian --quantization 0 --epochs --scale --lr --norm-clip 77 | python train.py --save-model --dataset mnist --model --mechanism gaussian --quantization 1 --epochs --scale --lr --norm-clip 78 | python train.py --save-model --dataset mnist --model --mechanism skellam --quantization 16 --epochs --scale --lr --norm-clip 79 | python train.py --save-model --dataset mnist --model --mechanism mvu --input-bits 9 --quantization 1 --beta 1 --epochs --epsilon --lr --norm-clip 80 | python train.py --save-model --dataset mnist --model --mechanism mvu_l2 --input-bits 1 --quantization 1 --beta 1 --epochs --epsilon --lr --norm-clip 81 | ``` 82 | To train on CIFAR-10, simply replace `--dataset mnist` by `--dataset cifar10`. See appendix in our paper for the full grid of hyperparameter values. 83 | 84 | ## Code Acknowledgements 85 | 86 | The majority of Privacy-Aware Compression is licensed under CC-BY-NC, however portions of the project are available under separate license terms: CVXPY and Opacus are licensed under the Apache 2.0 license; Kymatio is licensed under the BSD license; and Handcrafted-DP is licensed under the MIT license. 87 | -------------------------------------------------------------------------------- /compute_renyi_div.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | from mechanisms import * 10 | import numpy as np 11 | from tqdm import tqdm 12 | from utils import renyi_div_bound_lp 13 | import argparse 14 | 15 | if __name__ == "__main__": 16 | parser = argparse.ArgumentParser(description="Compute Renyi divergence for the MVU mechanism.") 17 | parser.add_argument( 18 | "--mechanism_folder", 19 | default="sweep_eps_budget_penalized_lam1.0e+02", 20 | type=str, 21 | help="folder containing saved optimal mechanisms", 22 | ) 23 | parser.add_argument( 24 | "--d", 25 | default=128, 26 | type=int, 27 | help="data dimensionality; must be a power of 2", 28 | ) 29 | parser.add_argument( 30 | "--Delta", 31 | default=0.5, 32 | type=float, 33 | help="input sensitivity Delta" 34 | ) 35 | parser.add_argument( 36 | "--epsilon", 37 | default=1, 38 | type=float, 39 | help="LDP epsilon", 40 | ) 41 | parser.add_argument( 42 | "--budget", 43 | default=3, 44 | type=int, 45 | help="budget for the MVU mechanism", 46 | ) 47 | parser.add_argument( 48 | "--input_bits", 49 | default=5, 50 | type=int, 51 | help="number of input bits for the MVU mechanism", 52 | ) 53 | parser.add_argument( 54 | "--dp_constraint", 55 | default="strict", 56 | type=str, 57 | choices=["strict", "metric-l1", "metric-l2"], 58 | help="type of DP constraint" 59 | ) 60 | args = parser.parse_args() 61 | 62 | alphas = np.array(list(np.linspace(1.1, 10.9, 99)) + list(range(11, 64))) 63 | savefile = os.path.join(args.mechanism_folder, 64 | f"mechanism_bin{args.input_bits}_bout{args.budget}_{args.dp_constraint}_eps{args.epsilon:.2f}.pkl") 65 | with open(savefile, "rb") as file: 66 | mechanism = pickle.load(file) 67 | renyi_div_bound = renyi_div_bound_lp(alphas, args.d, mechanism.P, args.Delta) 68 | output_file = f"renyi_div_bin{args.input_bits}_bout{args.budget}_{args.dp_constraint}_eps{args.epsilon:.2f}_Delta{args.Delta:.2f}.npz" 69 | np.savez(os.path.join(args.mechanism_folder, output_file), alphas=alphas, renyi_div_bound=renyi_div_bound) 70 | -------------------------------------------------------------------------------- /install.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | """ 10 | Installs dependencies for open-source users. 11 | """ 12 | 13 | import os 14 | import requests 15 | 16 | 17 | # files for RDP dependency: 18 | RDP_BASE = "https://raw.githubusercontent.com/tensorflow/privacy" 19 | RDP_COMMIT = "1ce8cd4032b06e8afa475747a105cfcb01c52ebe" 20 | RDP_FOLDER = "tensorflow_privacy/privacy/analysis" 21 | RDP_FILES = { 22 | "rdp_accountant.py": "rdp_accountant.py", 23 | "compute_dp_sgd_privacy.py": "dpsgd_privacy.py" 24 | } 25 | 26 | # files for ResNet dependency: 27 | RESNET_BASE = "https://raw.githubusercontent.com/akamaster/pytorch_resnet_cifar10" 28 | RESNET_COMMIT = "4e4f8da1ba2611dad2eedf8a23505c0fbd94b983" 29 | RESNET_FILES = { 30 | "resnet.py": "resnet.py", 31 | } 32 | 33 | PP_FILES = { 34 | "private_prediction.py": "private_prediction.py", 35 | } 36 | 37 | 38 | def main(): 39 | 40 | for filename in PP_FILES.values(): 41 | patch_filename = filename.replace(".py", ".patch") 42 | os.system(f"patch private_prediction/{filename} patches/{patch_filename}") 43 | 44 | pp_dir = "private_prediction" 45 | 46 | # download files needed for RDP accountant: 47 | for source, filename in RDP_FILES.items(): 48 | url = f"{RDP_BASE}/{RDP_COMMIT}/{RDP_FOLDER}/{source}" 49 | request = requests.get(url, allow_redirects=True) 50 | with open(os.path.join(pp_dir, filename), "wb") as f: 51 | print(f"writing file {filename}") 52 | f.write(request.content) 53 | 54 | # download files needed for ResNet: 55 | for source, filename in RESNET_FILES.items(): 56 | url = f"{RESNET_BASE}/{RESNET_COMMIT}/{source}" 57 | request = requests.get(url, allow_redirects=True) 58 | with open(os.path.join(pp_dir, filename), "wb") as f: 59 | print(f"writing file {filename}") 60 | f.write(request.content) 61 | 62 | # apply patches: 63 | for filename in RDP_FILES.values(): 64 | patch_filename = filename.replace(".py", ".patch") 65 | os.system(f"patch {pp_dir}/{filename} {pp_dir}/patches/{patch_filename}") 66 | for filename in RESNET_FILES.values(): 67 | patch_filename = filename.replace(".py", ".patch") 68 | os.system(f"patch {pp_dir}/{filename} {pp_dir}/patches/{patch_filename}") 69 | print("done.") 70 | 71 | 72 | if __name__ == "__main__": 73 | main() 74 | -------------------------------------------------------------------------------- /mean_estimation_multi.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | from mechanisms import * 10 | import numpy as np 11 | from tqdm import tqdm 12 | from opacus.accountants.analysis.rdp import get_privacy_spent 13 | import argparse 14 | from utils import optimal_scaling_mvu, optimal_scaling_skellam 15 | import torch.nn as nn 16 | 17 | import sys 18 | sys.path.append("private_prediction/") 19 | from util import binary_search 20 | from private_prediction import sensitivity_scale 21 | 22 | 23 | if __name__ == "__main__": 24 | parser = argparse.ArgumentParser(description="Run vector distributed mean estimation experiment.") 25 | parser.add_argument( 26 | "--save_folder", 27 | default="dme_results", 28 | type=str, 29 | help="folder in which to store results", 30 | ) 31 | parser.add_argument( 32 | "--mechanism_folder", 33 | default="sweep_eps_budget_penalized_lam1.0e+02", 34 | type=str, 35 | help="folder containing saved MVU mechanisms", 36 | ) 37 | parser.add_argument( 38 | "--trials", 39 | default=10, 40 | type=int, 41 | help="number of trials", 42 | ) 43 | parser.add_argument( 44 | "--num_samples", 45 | default=10000, 46 | type=int, 47 | help="number of samples", 48 | ) 49 | parser.add_argument( 50 | "--d", 51 | default=128, 52 | type=int, 53 | help="data dimensionality; must be a power of 2", 54 | ) 55 | parser.add_argument( 56 | "--norm_type", 57 | default="l1", 58 | choices=["l1", "l2"], 59 | type=str, 60 | help="generate synthetic data under L1 or L2 norm bound", 61 | ) 62 | parser.add_argument( 63 | "--epsilon", 64 | default=1, 65 | type=float, 66 | help="LDP epsilon", 67 | ) 68 | parser.add_argument( 69 | "--skellam_budget", 70 | default=16, 71 | type=int, 72 | help="budget for the Skellam mechanism", 73 | ) 74 | parser.add_argument( 75 | "--skellam_s", 76 | default=15, 77 | type=float, 78 | help="scaling factor for the Skellam mechanism", 79 | ) 80 | parser.add_argument( 81 | "--mvu_budget", 82 | default=16, 83 | type=int, 84 | help="budget for the MVU mechanism", 85 | ) 86 | parser.add_argument( 87 | "--mvu_input_bits", 88 | default=5, 89 | type=int, 90 | help="number of input bits for the MVU mechanism", 91 | ) 92 | parser.add_argument( 93 | "--dither_tol", 94 | default=0.1, 95 | type=float, 96 | help="failure probability for conditional dithering", 97 | ) 98 | args = parser.parse_args() 99 | os.makedirs(args.save_folder, exist_ok=True) 100 | 101 | if args.norm_type == "l1": 102 | p = 1 103 | # generate from uniform([0, 1]) then normalize 104 | xs = np.random.random((args.num_samples, args.d)) 105 | else: 106 | p = 2 107 | # generate from uniform over positive quadrant of unit sphere 108 | xs = np.absolute(np.random.normal(0, 1, (args.num_samples, args.d))) 109 | xs /= np.maximum(np.linalg.norm(xs, p, 1), 1)[:, None] 110 | 111 | # CLDP 112 | mechanism_cldp = CLDPMechanism(args.epsilon, args.d, 1, args.norm_type) 113 | 114 | # Skellam 115 | mechanism = SkellamMechanism(args.skellam_budget, args.d, 1, 1, args.skellam_s) 116 | skellam_scale = optimal_scaling_skellam(xs, mechanism, args.skellam_s, args.dither_tol, p) 117 | 118 | mus = np.power(10, np.linspace(-2, 2, 100)) 119 | orders = np.array(list(np.linspace(1.1, 10.9, 99)) + list(range(11, 64))) 120 | for mu in mus: 121 | mechanism_skellam = SkellamMechanism(args.skellam_budget, args.d, 1, mu, args.skellam_s, p=p) 122 | rdp_const = mechanism_skellam.renyi_div(orders) 123 | epsilon_opt, _ = get_privacy_spent(orders=orders, rdp=rdp_const, delta=(1/(args.num_samples+1))) 124 | if epsilon_opt < args.epsilon: 125 | print("Optimal mu = %.2f" % mu) 126 | break 127 | 128 | # MVU 129 | epsilon = 2 * args.epsilon if args.norm_type == "l1" else 4 * args.epsilon 130 | mechanism = MultinomialSamplingMechanism(args.mvu_budget, epsilon, args.mvu_input_bits, norm_bound=0.5, p=None) 131 | mvu_scale = optimal_scaling_mvu(xs, mechanism, args.dither_tol, p) 132 | savefile = os.path.join( 133 | args.mechanism_folder, f"mechanism_bin{args.mvu_input_bits}_bout{args.mvu_budget}_metric-{args.norm_type}_eps{epsilon:.2f}.pkl") 134 | with open(savefile, "rb") as file: 135 | mechanism_mvu = pickle.load(file) 136 | mechanism_mvu.P /= mechanism_mvu.P.sum(1)[:, None] 137 | mechanism_mvu.norm_bound = 0.5 138 | mechanism_mvu.p = p 139 | 140 | # MVU mechanism for approximate DP 141 | epsilon = 2 * args.epsilon if args.norm_type == "l1" else 0.5 * args.epsilon 142 | savefile = os.path.join( 143 | args.mechanism_folder, f"mechanism_bin{args.mvu_input_bits}_bout{args.mvu_budget}_metric-l1_eps{epsilon:.2f}.pkl") 144 | with open(savefile, "rb") as file: 145 | mechanism_mvu_approx = pickle.load(file) 146 | mechanism_mvu_approx.P /= mechanism_mvu_approx.P.sum(1)[:, None] 147 | mechanism_mvu_approx.norm_bound = 0.5 148 | mechanism_mvu_approx.p = p 149 | 150 | squared_error_cldp = np.zeros(args.trials) 151 | squared_error_skellam = np.zeros(args.trials) 152 | squared_error_mvu = np.zeros(args.trials) 153 | squared_error_mvu_approx = np.zeros(args.trials) 154 | squared_error_baseline = np.zeros(args.trials) 155 | 156 | for k in tqdm(range(args.trials)): 157 | if args.norm_type == "l1": 158 | # generate from uniform([0, 1]) then normalize 159 | xs = np.random.random((args.num_samples, args.d)) 160 | else: 161 | # generate from uniform over positive quadrant of unit sphere 162 | xs = np.absolute(np.random.normal(0, 1, (args.num_samples, args.d))) 163 | xs /= np.maximum(np.linalg.norm(xs, p, 1), 1)[:, None] 164 | mean = xs.mean(0) 165 | 166 | mean_cldp = np.zeros(xs.shape[1]) 167 | for i in range(args.num_samples): 168 | x = xs[i] 169 | result = mechanism_cldp.decode(mechanism_cldp.privatize(x)) 170 | mean_cldp += result / args.num_samples 171 | squared_error_cldp[k] = np.power(mean - mean_cldp, 2).mean() 172 | 173 | mean_skellam = np.zeros(xs.shape[1]) 174 | for i in range(args.num_samples): 175 | x = skellam_scale * xs[i] 176 | output = mechanism_skellam.privatize(x) 177 | mean_skellam += output 178 | mean_skellam = np.mod(mean_skellam - mechanism_skellam.clip_min, mechanism_skellam.clip_max - mechanism_skellam.clip_min) + mechanism_skellam.clip_min 179 | mean_skellam = mechanism_skellam.decode(mean_skellam) / (skellam_scale * args.num_samples) 180 | squared_error_skellam[k] = np.power(mean - mean_skellam, 2).mean() 181 | 182 | mean_mvu = np.zeros(xs.shape[1]) 183 | prepro = lambda z: mvu_scale * z 184 | prepro_inv = lambda z: z / mvu_scale 185 | for i in range(args.num_samples): 186 | x = prepro(xs[i]) 187 | x = np.clip((x + 1) / 2, 0, 1) 188 | result = 2 * mechanism_mvu.decode(mechanism_mvu.privatize(x)) - 1 189 | mean_mvu += prepro_inv(result) / args.num_samples 190 | squared_error_mvu[k] = np.power(mean - mean_mvu, 2).mean() 191 | 192 | mean_mvu_approx = np.zeros(xs.shape[1]) 193 | for i in range(args.num_samples): 194 | x = prepro(xs[i]) 195 | x = np.clip((x + 1) / 2, 0, 1) 196 | result = 2 * mechanism_mvu_approx.decode(mechanism_mvu_approx.privatize(x)) - 1 197 | mean_mvu_approx += prepro_inv(result) / args.num_samples 198 | squared_error_mvu_approx[k] = np.power(mean - mean_mvu_approx, 2).mean() 199 | 200 | if args.norm_type == "l1": 201 | mean_baseline = (xs + np.random.laplace(0, 1 / args.epsilon, size=xs.shape)).mean(0) 202 | else: 203 | std = 1 / sensitivity_scale(args.epsilon, 1/(args.num_samples+1), None, None, None, 204 | "advanced_gaussian", chaudhuri=False) 205 | mean_baseline = (xs + np.random.normal(0, std, size=xs.shape)).mean(0) 206 | squared_error_baseline[k] = np.power(mean - mean_baseline, 2).mean() 207 | 208 | savefile = "%s/dme_multi_%s_d_%d_samples_%d_eps_%.2f_skellam_%d_%.2f_mvu_%d_%d.npz" % ( 209 | args.save_folder, args.norm_type, args.d, args.num_samples, args.epsilon, args.skellam_budget, args.skellam_s, args.mvu_budget, args.mvu_input_bits) 210 | np.savez(savefile, squared_error_cldp=squared_error_cldp, squared_error_skellam=squared_error_skellam, 211 | squared_error_mvu=squared_error_mvu, squared_error_mvu_approx=squared_error_mvu_approx, 212 | squared_error_baseline=squared_error_baseline) 213 | -------------------------------------------------------------------------------- /mean_estimation_single.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | from mechanisms import * 10 | import numpy as np 11 | from tqdm import tqdm 12 | import argparse 13 | import seaborn as sns 14 | import matplotlib.pyplot as plt 15 | 16 | 17 | def compute_error(xs, mechanism, num_trials): 18 | squared_errors = np.zeros(xs.shape) 19 | for i in tqdm(range(len(xs))): 20 | x = xs[i] 21 | inputs = x * np.ones(num_trials) 22 | result = mechanism.decode(mechanism.privatize(inputs)) 23 | squared_errors[i] = np.power(result - x, 2).mean() 24 | return squared_errors 25 | 26 | 27 | def compute_error_1d(xs, mechanism, num_trials): 28 | means = np.zeros(xs.shape) 29 | squared_errors = np.zeros(xs.shape) 30 | for i in tqdm(range(len(xs))): 31 | x = np.array([xs[i]]) 32 | for _ in range(num_trials): 33 | result = mechanism.decode(mechanism.privatize(x)) 34 | means[i] += result / num_trials 35 | squared_errors[i] += np.power(result - x, 2) / num_trials 36 | return means, squared_errors 37 | 38 | 39 | if __name__ == "__main__": 40 | parser = argparse.ArgumentParser(description="Run scalar distributed mean estimation experiment.") 41 | parser.add_argument( 42 | "--save_folder", 43 | default="figures", 44 | type=str, 45 | help="folder in which to save the figures", 46 | ) 47 | parser.add_argument( 48 | "--mechanism_folder", 49 | default="sweep_eps_budget_tr", 50 | type=str, 51 | help="folder containing saved MVU mechanisms", 52 | ) 53 | parser.add_argument( 54 | "--num_samples", 55 | default=int(1e5), 56 | type=int, 57 | help="number of samples", 58 | ) 59 | parser.add_argument( 60 | "--epsilon", 61 | default=1, 62 | type=float, 63 | help="LDP epsilon", 64 | ) 65 | parser.add_argument( 66 | "--budget", 67 | default=3, 68 | type=int, 69 | help="budget for the MVU mechanism", 70 | ) 71 | parser.add_argument( 72 | "--input_bits", 73 | default=3, 74 | type=int, 75 | help="number of input bits for the MVU mechanism", 76 | ) 77 | args = parser.parse_args() 78 | os.makedirs(args.save_folder, exist_ok=True) 79 | 80 | xs = np.linspace(-1, 1, 50) 81 | xs_normalized = (xs + 1) / 2 82 | 83 | print("Running RR") 84 | mechanism = RandomizedResponseMechanism(args.budget, args.epsilon) 85 | squared_errors_rr = 4 * compute_error(xs_normalized, mechanism, args.num_samples) 86 | 87 | print("Running Generalized RR") 88 | mechanism = GeneralizedRRMechanism(args.budget, args.epsilon, args.budget) 89 | squared_errors_rappor = 4 * compute_error(xs_normalized, mechanism, args.num_samples) 90 | 91 | print("Running CLDP") 92 | mechanism = CLDPMechanism(args.epsilon, 1, 1, "linf") 93 | squared_errors_cldp = compute_error_1d(xs, mechanism, args.num_samples)[1] 94 | 95 | print("Running MVU b=%d" % args.budget) 96 | savefile = os.path.join( 97 | args.mechanism_folder, f"mechanism_bin{args.input_bits}_bout{args.budget}_strict_eps{args.epsilon:.2f}.pkl") 98 | if os.path.exists(savefile): 99 | with open(savefile, "rb") as file: 100 | mechanism = pickle.load(file) 101 | else: 102 | mechanism = MVUMechanism(args.budget, args.epsilon, args.input_bits, method="trust-region", init_method="uniform") 103 | squared_errors_mvu = 4 * compute_error(xs_normalized, mechanism, args.num_samples) 104 | 105 | print("Running MVU b=1") 106 | savefile = os.path.join( 107 | args.mechanism_folder, f"mechanism_bin{args.input_bits}_bout1_strict_eps{args.epsilon:.2f}.pkl") 108 | if os.path.exists(savefile): 109 | with open(savefile, "rb") as file: 110 | mechanism = pickle.load(file) 111 | else: 112 | mechanism = MVUMechanism(1, args.epsilon, args.input_bits, method="trust-region", init_method="uniform") 113 | squared_errors_mvu_1bit = 4 * compute_error(xs_normalized, mechanism, args.num_samples) 114 | 115 | plt.figure(figsize=(8,5)) 116 | colors = sns.color_palette("deep") 117 | plt.plot(xs, squared_errors_rr, label='Bitwise RR', color=colors[6], linewidth=3) 118 | plt.plot(xs, squared_errors_rappor, label='Generalized RR', color=colors[3], linewidth=3) 119 | plt.plot(xs, squared_errors_cldp, label='CLDP', color=colors[0], linewidth=3) 120 | plt.plot(xs, squared_errors_mvu_1bit, label='MVU ($b=1$)', color='lightgreen', linewidth=3) 121 | plt.plot(xs, squared_errors_mvu, label='MVU ($b=3$)', color=colors[2], linewidth=3) 122 | plt.plot(xs, np.ones(xs.shape) * 2 / args.epsilon**2, label='Laplace', color='k', linestyle='--', linewidth=3) 123 | 124 | plt.xlabel('$x$', fontsize=20) 125 | plt.xticks([-1, -0.5, 0, 0.5, 1], fontsize=20) 126 | plt.ylabel('Variance', fontsize=20) 127 | plt.yticks(fontsize=20) 128 | plt.grid('on') 129 | plt.legend(loc='upper right', fontsize=20) 130 | plt.savefig("%s/dme_single_eps_%.2f.pdf" % (args.save_folder, args.epsilon), bbox_inches="tight") 131 | -------------------------------------------------------------------------------- /mechanisms.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | import math 10 | import random 11 | import numpy as np 12 | from abc import abstractmethod 13 | import scipy 14 | from scipy.stats import skellam 15 | import scipy.optimize as optimize 16 | from scipy import sparse, optimize 17 | from scipy.special import gamma, softmax 18 | import cvxpy 19 | import pdb 20 | from scipy.sparse.csr import csr_matrix 21 | from tqdm import tqdm 22 | import os 23 | from datetime import datetime 24 | import pickle 25 | from utils import optimal_scaling_integer, FWHTRandomProjector 26 | import time 27 | from opacus.accountants import RDPAccountant 28 | from typing import Callable, List, Optional, Union, Tuple 29 | from opacus.accountants.analysis import rdp as privacy_analysis 30 | 31 | 32 | class CLDPMechanism: 33 | ''' 34 | LDP mechanisms from https://arxiv.org/pdf/2008.07180.pdf. 35 | ''' 36 | 37 | def __init__(self, epsilon, d, norm_bound, norm_type): 38 | self.epsilon = epsilon 39 | if norm_type == "l1": 40 | self.d = int(math.pow(2, math.ceil(math.log2(d)))) 41 | else: 42 | self.d = d 43 | self.norm_bound = norm_bound 44 | self.norm_type = norm_type 45 | 46 | def privatize_l1(self, x): 47 | assert np.linalg.norm(x, 1) <= self.norm_bound + 1e-8 48 | assert len(x) <= self.d 49 | z = np.zeros(self.d) 50 | z[:len(x)] = x 51 | H = scipy.linalg.hadamard(self.d) 52 | z = H @ z / math.sqrt(self.d) 53 | C = (math.exp(self.epsilon) - 1) / (math.exp(self.epsilon) + 1) 54 | idx = random.randint(0, self.d - 1) 55 | U = np.sign(0.5 + C * math.sqrt(self.d) * z[idx] / (2 * self.norm_bound) - random.random()) 56 | return (idx, U) 57 | 58 | def decode_l1(self, z): 59 | C_inv = (math.exp(self.epsilon) + 1) / (math.exp(self.epsilon) - 1) 60 | H = scipy.linalg.hadamard(self.d) 61 | return z[1] * self.norm_bound * C_inv * H[:, z[0]] 62 | 63 | def quantize_l2(self, x): 64 | norm = np.linalg.norm(x, 1) 65 | sign = np.sign((1 + norm) / (2 * self.norm_bound * math.sqrt(self.d)) - random.random()) 66 | x = sign * x / norm 67 | if len(x) > 1: 68 | y = np.random.multinomial(1, np.absolute(x), (self.d,)) 69 | else: 70 | y = np.ones((self.d, 1)).astype(int) 71 | return y, np.sign(x[np.argmax(y, 1)]) 72 | 73 | def privatize_l2(self, x): 74 | assert np.linalg.norm(x, 2) <= self.norm_bound + 1e-8 75 | d = len(x) 76 | norm = np.linalg.norm(x, 2) 77 | sign = np.sign(0.5 + norm / (2 * self.norm_bound) - random.random()) 78 | x = sign * self.norm_bound * x / norm 79 | U_sign = np.sign(math.exp(self.epsilon) / (math.exp(self.epsilon) + 1) - random.random()) 80 | C_inv = (math.exp(self.epsilon) + 1) / (math.exp(self.epsilon) - 1) 81 | M = self.norm_bound * d * math.sqrt(math.pi) * C_inv * gamma((d-1)/2 + 1) / (2 * gamma(d/2 + 1)) 82 | while True: 83 | z = np.random.normal(0, 1, (d,)) 84 | z = M * z / np.linalg.norm(z) 85 | if sum(z * x) * U_sign > 0: 86 | break 87 | return self.quantize_l2(z) 88 | 89 | def decode_l2(self, z): 90 | return self.norm_bound * (z[0] * z[1][:, None]).sum(0) 91 | 92 | def privatize_linf(self, x): 93 | assert np.absolute(x).max() <= self.norm_bound + 1e-8 94 | C = (math.exp(self.epsilon) - 1) / (math.exp(self.epsilon) + 1) 95 | idx = random.randint(0, len(x) - 1) 96 | U = np.sign(0.5 + C * x[idx] / (2 * self.norm_bound) - random.random()) 97 | return (idx, U) 98 | 99 | def decode_linf(self, z): 100 | C_inv = (math.exp(self.epsilon) + 1) / (math.exp(self.epsilon) - 1) 101 | e = np.zeros(self.d) 102 | e[z[0]] = 1 103 | return z[1] * self.norm_bound * self.d * C_inv * e 104 | 105 | def privatize(self, x): 106 | if self.norm_type == "l1": 107 | return self.privatize_l1(x) 108 | elif self.norm_type == "l2": 109 | return self.privatize_l2(x) 110 | elif self.norm_type == "linf": 111 | return self.privatize_linf(x) 112 | else: 113 | raise RuntimeError("Unsupported norm type: " + str(self.norm_type)) 114 | 115 | def decode(self, x): 116 | if self.norm_type == "l1": 117 | return self.decode_l1(x) 118 | elif self.norm_type == "l2": 119 | return self.decode_l2(x) 120 | elif self.norm_type == "linf": 121 | return self.decode_linf(x) 122 | else: 123 | raise RuntimeError("Unsupported norm type: " + str(self.norm_type)) 124 | 125 | 126 | class SkellamMechanism: 127 | ''' 128 | Skellam mechanism from https://arxiv.org/pdf/2110.04995.pdf. 129 | ''' 130 | 131 | def __init__(self, budget, d, norm_bound, mu, num_clients=1): 132 | self.budget = budget 133 | self.d = d 134 | self.expanded_d = int(math.pow(2, math.ceil(math.log2(d)))) 135 | self.norm_bound = norm_bound 136 | self.mu = mu 137 | self.s = self.compute_s(num_clients) 138 | print("s = %.2f" % self.s) 139 | self.scale = optimal_scaling_integer(self.expanded_d, self.s * norm_bound, math.exp(-0.5), tol=1e-3) 140 | if self.scale == 0: 141 | raise RuntimeError("Did not find suitable scale factor; try increasing communication budget") 142 | self.clip_min = -int(math.pow(2, budget - 1)) 143 | self.clip_max = int(math.pow(2, budget - 1)) - 1 144 | self.projector = FWHTRandomProjector(self.d, self.expanded_d) 145 | return 146 | 147 | def compute_s(self, num_clients, k=3, rho=1, DIV_EPSILON=1e-22): 148 | """ 149 | Adapted from https://github.com/google-research/federated/blob/master/distributed_dp/accounting_utils.py 150 | """ 151 | def mod_min(gamma): 152 | var = rho / self.d * (num_clients * self.norm_bound)**2 153 | var += (gamma**2 / 4 + self.mu) * num_clients 154 | return k * math.sqrt(var) 155 | 156 | def gamma_opt_fn(gamma): 157 | return (math.pow(2, self.budget) - 2 * mod_min(gamma) / (gamma + DIV_EPSILON))**2 158 | 159 | gamma_result = optimize.minimize_scalar(gamma_opt_fn) 160 | if not gamma_result.success: 161 | raise ValueError('Cannot compute scaling factor.') 162 | return 1. / gamma_result.x 163 | 164 | def renyi_div(self, alphas, l1_norm_bound=None, l2_norm_bound=None): 165 | """ 166 | Computes Renyi divergence of the Skellam mechanism. 167 | """ 168 | if l2_norm_bound is None: 169 | l2_norm_bound = self.norm_bound 170 | if l1_norm_bound is None: 171 | l1_norm_bound = self.norm_bound * min(math.sqrt(self.expanded_d), self.norm_bound) 172 | epsilons = np.zeros(alphas.shape) 173 | B1 = 3 * l1_norm_bound / (2 * self.s ** 3 * self.mu ** 2) 174 | B2 = 3 * l1_norm_bound / (2 * self.s * self.mu) 175 | for i in range(len(alphas)): 176 | alpha = alphas[i] 177 | epsilon = alpha * self.norm_bound ** 2 / (2 * self.mu) 178 | B3 = (2 * alpha - 1) * self.norm_bound ** 2 / (4 * self.s ** 2 * self.mu ** 2) 179 | epsilons[i] = epsilon + min(B1 + B3, B2) 180 | return epsilons 181 | 182 | def dither(self, x): 183 | k = np.floor(x) 184 | prob = 1 - (x - k) 185 | while True: 186 | output = k + (np.random.random(k.shape) > prob) 187 | if np.linalg.norm(output, 2) <= self.s * self.norm_bound: 188 | break 189 | return output.astype(int) 190 | 191 | def privatize(self, x): 192 | assert np.all(np.linalg.norm(x, 2, 1) <= self.norm_bound + 1e-4) 193 | assert x.shape[1] == self.d 194 | z = np.zeros((x.shape[0], self.expanded_d)) 195 | for i in range(x.shape[0]): 196 | z[i] = self.dither(self.projector.project(self.s * x[i])) 197 | z += skellam.rvs(self.s**2 * self.mu, self.s**2 * self.mu, size=z.shape) 198 | z = np.mod(z - self.clip_min, self.clip_max - self.clip_min) + self.clip_min 199 | return z 200 | 201 | def decode(self, z): 202 | x = np.zeros((z.shape[0], self.d)) 203 | for i in range(z.shape[0]): 204 | x[i] = self.projector.inverse(z[i].astype(float)) / self.s 205 | return x 206 | 207 | 208 | class SkellamAccountant(RDPAccountant): 209 | 210 | def __init__(self, orders, renyi_div_bounds): 211 | super().__init__() 212 | self.alphas = orders 213 | self.renyi_div_bounds = renyi_div_bounds 214 | 215 | def get_privacy_spent( 216 | self, *, delta: float, alphas: Optional[List[Union[float, int]]] = None 217 | ) -> Tuple[float, float]: 218 | if not self.history: 219 | return 0, 0 220 | 221 | # MVU accountant does not yet support subsampling and different noise multipliers 222 | rdp = sum([self.renyi_div_bounds * num_steps for (_, _, num_steps) in self.history]) 223 | eps, best_alpha = privacy_analysis.get_privacy_spent( 224 | orders=self.alphas, rdp=rdp, delta=delta 225 | ) 226 | return float(eps), float(best_alpha) 227 | 228 | 229 | class CompressedMechanism: 230 | 231 | def __init__(self, budget, epsilon): 232 | self.budget = budget 233 | self.epsilon = epsilon 234 | return 235 | 236 | def dither(self, x, b, p=None): 237 | """ 238 | Given x in [0,1], return a randomized dithered output in {0, 1, ..., 2^b - 1}. 239 | """ 240 | assert np.all(x >= 0) and np.all(x <= 1) 241 | B = 2 ** b 242 | k = np.floor((B-1) * x) 243 | prob = 1 - (B-1) * (x - k/(B-1)) 244 | k += np.random.random(k.shape) > prob 245 | return k.astype(int) 246 | 247 | @abstractmethod 248 | def privatize(self, x): 249 | """ 250 | Privatizes a vector of values in [0,1] to binary vectors. 251 | """ 252 | return 253 | 254 | @abstractmethod 255 | def decode(self, l): 256 | """ 257 | Decodes binary vectors to an array of real values. 258 | """ 259 | return 260 | 261 | 262 | class RandomizedResponseMechanism(CompressedMechanism): 263 | 264 | def _privatize_bit(self, x, epsilon): 265 | """ 266 | Privatizes a vector of bits using the binary randomized response mechanism. 267 | """ 268 | assert np.all(np.logical_or(x == 0, x == 1)) 269 | prob = 1 / (1 + math.exp(-epsilon)) 270 | mask = np.random.random(x.shape) > prob 271 | z = np.logical_xor(mask, x).astype(int) 272 | return z 273 | 274 | def binary_repr(self, x): 275 | """ 276 | Converts an array of integers to a 2D array of bits using binary representation. 277 | """ 278 | l = [np.fromiter(map(int, np.binary_repr(a, width=self.budget)), int) for a in x] 279 | return np.stack(l, 0) 280 | 281 | def int_repr(self, l): 282 | """ 283 | Converts a 2D array of bits into an array of integers using binary representation. 284 | """ 285 | powers = np.power(2, np.arange(self.budget-1, -0.5, -1)) 286 | return l.dot(powers) 287 | 288 | def privatize(self, x): 289 | z = self.dither(x, self.budget) 290 | l = self.binary_repr(z) 291 | l = self._privatize_bit(l, float(self.epsilon/self.budget)) 292 | return l 293 | 294 | def decode(self, l): 295 | assert l.shape[1] == self.budget 296 | a_0 = -1 / (math.exp(self.epsilon/self.budget) - 1) 297 | a_1 = math.exp(self.epsilon/self.budget) / (math.exp(self.epsilon/self.budget) - 1) 298 | l = a_0 + l * (a_1 - a_0) 299 | return self.int_repr(l) / (2**self.budget - 1) 300 | 301 | 302 | class MultinomialSamplingMechanism(CompressedMechanism): 303 | 304 | def __init__(self, budget, epsilon, input_bits, norm_bound, p, **kwargs): 305 | """ 306 | Parent class that supports sampling from a 2^budget-dimensional distribution defined by 307 | a sampling probability matrix P and an output vector alpha. 308 | 309 | Arguments: 310 | budget - Number of bits in the output. 311 | epsilon - DP/metric-DP parameter epsilon. 312 | input_bits - Number of bits in the quantized input. 313 | norm_bound - A priori bound on the norm of the input before quantization; ignored if p=None. 314 | p - Which p-norm to use for the norm bound parameter. 315 | """ 316 | super().__init__(budget, epsilon) 317 | self.input_bits = input_bits 318 | self.norm_bound = norm_bound 319 | self.p = p 320 | result = self.optimize(**kwargs) 321 | if result is not None: 322 | self.P, self.alpha = result[0], result[1] 323 | return 324 | 325 | def dither(self, x, b): 326 | """ 327 | Dithers x coordinate-wise to a grid of size 2^b. 328 | If self.p is set, perform rejection sampling until dithered vector does not exceed self.norm_bound. 329 | """ 330 | assert np.all(x >= 0) and np.all(x <= 1) 331 | B = 2 ** b 332 | k = np.floor((B-1) * x) 333 | prob = 1 - (B-1) * (x - k/(B-1)) 334 | while True: 335 | output = k + (np.random.random(k.shape) > prob) 336 | if self.p is None or np.linalg.norm(output / (B-1) - 0.5, self.p) <= self.norm_bound: 337 | break 338 | return output.astype(int) 339 | 340 | @abstractmethod 341 | def optimize(self, **kwargs): 342 | """ 343 | Optimizes self.P and self.alpha for multinomial sampling. 344 | """ 345 | return 346 | 347 | def privatize(self, x): 348 | z = self.dither(x, self.input_bits) 349 | B = 2**self.budget 350 | range_B = np.arange(B).astype(int) 351 | z = np.array([np.random.choice(range_B, p=self.P[a]) for a in z]) 352 | return z 353 | 354 | def decode(self, z): 355 | assert z.min() >= 0 and z.max() < 2**self.budget 356 | return self.alpha[z.astype(int)] 357 | 358 | def mse_and_bias_squared(self, P=None, alpha=None): 359 | """ 360 | Evaluate MSE loss and bias squared. 361 | """ 362 | if P is None and alpha is None: 363 | P = self.P 364 | alpha = self.alpha 365 | B = 2 ** self.input_bits 366 | target = np.arange(0, 1+1/B, 1/(B-1)) 367 | mse_loss = (P * np.power(target[:, None] - alpha[None, :], 2)).sum(1).mean() 368 | bias_sq = np.power(P @ alpha - target, 2).mean() 369 | return mse_loss, bias_sq 370 | 371 | 372 | class RAPPORMechanism(MultinomialSamplingMechanism): 373 | 374 | def __init__(self, budget, epsilon, input_bits, norm_bound=0.5, p=None, **kwargs): 375 | super().__init__(budget, epsilon, budget, norm_bound, p, **kwargs) # ignores input bits 376 | return 377 | 378 | def optimize(self): 379 | B = 2**self.budget 380 | prob = B / (B + math.exp(self.epsilon) - 1) 381 | P = prob / B * np.ones((B, B)) + (1 - prob) * np.eye(B) 382 | target = np.arange(0, 1+1/B, 1/(B-1)) 383 | alpha = np.linalg.solve(P, target) 384 | return P, alpha 385 | 386 | 387 | class MVUMechanism(MultinomialSamplingMechanism): 388 | 389 | def __init__(self, budget, epsilon, input_bits, norm_bound=0.5, p=None, **kwargs): 390 | super().__init__(budget, epsilon, input_bits, norm_bound, p, **kwargs) 391 | return 392 | 393 | # ================================== 394 | # Functions used by multiple methods 395 | # ================================== 396 | def _get_dp_constraint_matrix(self, dp_constraint, sparsity=0): 397 | """ 398 | Returns a sparse matrix D with shape (B*B*(B-1), B*B) such that D @ p 399 | corresponds to the left-hand side of the DP constraints on P, 400 | where p = P.reshape(B*B) and B = 2 ** self.budget. 401 | 402 | The final constriants are D @ p <= 0. 403 | 404 | Each row of D bounds the probability ratio between P_{i,j} and P_{i',j} 405 | by e^epsilon for i != i'. For metric DP, the ratio is e^{epsilon * abs(i - i')}. 406 | """ 407 | print(f"sparsity = {sparsity}") 408 | B_in = 2 ** self.input_bits 409 | B_out = 2 ** self.budget 410 | assert sparsity >= 0 and sparsity < B_in 411 | if sparsity == 0: 412 | if dp_constraint == "strict": 413 | data = -math.exp(self.epsilon) * np.ones(B_in*(B_in-1)) 414 | elif dp_constraint == "metric-l1" or dp_constraint == "metric-l2": 415 | data = np.absolute(np.arange(0, B_in)[:, None] - np.arange(0, B_in)[None, :]).reshape(B_in*B_in,) / (B_in-1) 416 | if dp_constraint == "metric-l2": 417 | data = np.power(data, 2) 418 | data = -np.exp(self.epsilon * data[data>0]) 419 | else: 420 | raise RuntimeError("Unknown DP constraint: " + str(dp_constraint)) 421 | row_indices = np.arange(0, B_in*(B_in-1)).astype(int) 422 | col_indices = np.floor(np.arange(0, B_in*(B_in-1)) / (B_in-1)).astype(int) 423 | coeffs = sparse.csr_matrix((data, (row_indices, col_indices)), shape=(B_in*(B_in-1), B_in)) 424 | D = sparse.kron(coeffs, sparse.eye(B_out)) # matrix of -e^{epsilons}s at all (i',j) positions 425 | E = sparse.csr_matrix( 426 | (np.ones(B_in), (np.arange(0, B_in*B_in, B_in+1).astype(int), np.arange(0, B_in).astype(int))), shape=(B_in*B_in, B_in)) 427 | F = sparse.kron(np.ones(B_in)[:, None], sparse.eye(B_in*B_out)) - sparse.kron(E, sparse.eye(B_out)) 428 | D += F[F.getnnz(1)>0] # add to D a matrix of 1s at all (i,j) positions 429 | else: 430 | D = [] 431 | for s in range(1, sparsity + 1): 432 | E = np.concatenate([np.eye(B_in-s), np.zeros((B_in-s, s))], 1) 433 | F = np.concatenate([np.zeros((B_in-s, s)), np.eye(B_in-s)], 1) 434 | if dp_constraint == "strict": 435 | epsilon = self.epsilon 436 | elif dp_constraint == "metric-l1": 437 | epsilon = self.epsilon * s / (B_in - 1) 438 | elif dp_constraint == "metric-l2": 439 | epsilon = self.epsilon * s**2 / (B_in - 1) ** 2 440 | else: 441 | raise RuntimeError("Unknown DP constraint: " + str(dp_constraint)) 442 | D.append(E - math.exp(epsilon) * F) 443 | D.append(F - math.exp(epsilon) * E) 444 | D = sparse.kron(sparse.csr_matrix(np.concatenate(D, 0)), sparse.eye(B_out)) 445 | return D 446 | 447 | def _get_row_stochastic_constraint_matrix(self): 448 | """ 449 | Returns a sparse matrix R with shape (B, B*B) such that R @ p corresponds to the 450 | left-hand side of the row-stochastic constraints, where p = P.reshape(B*B). 451 | 452 | The final constraints are R @ p == 1. 453 | """ 454 | B_in = 2 ** self.input_bits 455 | B_out = 2 ** self.budget 456 | return sparse.kron(sparse.eye(B_in), np.ones(B_out)) 457 | 458 | def _get_lp_costs(self, alpha, objective="squared"): 459 | """ 460 | When fixing alpha and solving for P, the design problem is a linear program in the entries of P. 461 | This function returns the cost function values in a vectorized form. 462 | 463 | The overall MSE objective is $\sum_{i,j} P[i,j] * (i/(B-1) - alpha[j])^2$. 464 | This function returns a vector of length B*B with the coefficient $(i/(B-1) - alpha[j])^2$ in 465 | the corresponding position. 466 | """ 467 | B_in = 2 ** self.input_bits 468 | target = np.arange(0, 1+1/B_in, 1/(B_in-1)) 469 | if objective == "squared": 470 | c = np.power(target[:, None] - alpha[None, :], 2).flatten() 471 | elif objective == "absolute": 472 | c = np.absolute(target[:, None] - alpha[None, :]).flatten() 473 | else: 474 | raise RuntimeError("Unknown objective: " + str(objective)) 475 | return c 476 | 477 | # =================== 478 | # Penalized LP method 479 | # =================== 480 | def _optimize_penalized_lp(self, objective="squared", dp_constraint="strict", sparsity=0, lam=0, num_iters=1, verbose=False): 481 | """ 482 | Estimate the optimal design by alternating between: 483 | a) Fixing alpha and solving for P, where the unbiased constraints are incorporated as a penalty in the objective, 484 | b) Fixing P and solving for alpha in terms of the linear system corresponding to unbiased constraints. 485 | """ 486 | B_in = 2 ** self.input_bits 487 | B_out = 2 ** self.budget 488 | alpha = np.arange(0, 1+1/B_out, 1/(B_out-1)) 489 | target = np.arange(0, 1+1/B_in, 1/(B_in-1)) 490 | num_var = B_in * B_out 491 | 492 | # Row-stochastic equality constraints: A_eq @ p == b_eq 493 | A_eq = self._get_row_stochastic_constraint_matrix() 494 | b_eq = np.ones(B_in) 495 | # Enforce symmetric P if input_bits is 1 496 | if self.input_bits == 1: 497 | A_sym = np.concatenate([np.eye(B_out), -np.fliplr(np.eye(B_out))], 1) 498 | A_sym = sparse.csr_matrix(A_sym) 499 | b_sym = np.zeros(B_out) 500 | A_eq = sparse.vstack([A_eq, A_sym]) 501 | b_eq = np.concatenate([b_eq, b_sym], 0) 502 | 503 | # DP inequality constraints: A_ineq @ p <= 0 504 | A_ineq = self._get_dp_constraint_matrix(dp_constraint, sparsity=sparsity) 505 | b_ineq = np.zeros(A_ineq.shape[0]) 506 | P = None 507 | 508 | for l in range(num_iters): 509 | 510 | c = self._get_lp_costs(alpha, objective) 511 | 512 | # Coefficients to implement the unbiased constraint as a quadratic penalty 513 | beta = np.kron(target[:, None], alpha[None, :]).reshape(1, -1)[0] 514 | Q = np.kron(np.eye(B_in), np.kron(alpha[:, None], alpha[None, :])) 515 | Q += 1e-6 * np.eye(Q.shape[0]) 516 | 517 | # Define and solve the CVXPY problem for P 518 | x = cvxpy.Variable(num_var, nonneg=True) 519 | if P is not None: 520 | x.value = P.flatten() 521 | if lam > 0: 522 | obj = cvxpy.quad_form(x, lam * Q) - 2 * lam * beta.T @ x + c.T @ x 523 | else: 524 | obj = c.T @ x 525 | prob = cvxpy.Problem(cvxpy.Minimize(obj), [A_ineq @ x <= b_ineq, A_eq @ x == b_eq]) 526 | try: 527 | prob.solve(solver="ECOS", max_iters=1000, warm_start=(P is not None), verbose=verbose) 528 | except (cvxpy.error.SolverError, scipy.sparse.linalg.eigen.arpack.ArpackNoConvergence): 529 | # Hack to continue even if solver fails to converge 530 | pass 531 | P = np.resize(x.value, (B_in, B_out)) 532 | 533 | if lam > 0: 534 | # Define and solve the CVXPY problem for alpha 535 | x = cvxpy.Variable(B_out, nonneg=False) 536 | x.value = alpha 537 | obj = cvxpy.quad_form(x, lam * P.T @ P + np.diag(P.sum(0))) - 2 * (1 + lam) * (target @ P) @ x 538 | prob = cvxpy.Problem(cvxpy.Minimize(obj)) 539 | prob.solve(solver="ECOS", max_iters=1000, warm_start=True, verbose=verbose) 540 | alpha = x.value 541 | else: 542 | # Obtain alpha by solving the linear system P @ alpha = target 543 | H = P.transpose().dot(P) + 1e-6 * np.eye(B_out) 544 | alpha = np.linalg.inv(H).dot(P.transpose().dot(target)) 545 | 546 | mse_loss, bias_sq = self.mse_and_bias_squared(P, alpha) 547 | if verbose: 548 | print("Iteration %d: MSE loss = %.4f, squared bias = %.4f" % ( 549 | l+1, mse_loss, bias_sq)) 550 | 551 | return P, alpha 552 | 553 | # =========================================== 554 | # Alternative method and supporting functions 555 | # =========================================== 556 | def _solve_lp_for_P(self, alpha, dp_constraint, sparsity, verbose=False): 557 | B_in = 2 ** self.input_bits 558 | B_out = 2 ** self.budget 559 | assert len(alpha) == B_out 560 | p = cvxpy.Variable(B_in * B_out, nonneg=True) 561 | u = np.arange(0, 1+1/B_in, 1/(B_in-1)) 562 | # Cost matrix as a function of alpha 563 | c = self._get_lp_costs(alpha) 564 | objective = cvxpy.Minimize(c.T @ p) 565 | # DP constraints 566 | D = self._get_dp_constraint_matrix(dp_constraint, sparsity) 567 | # Unbiased constraints 568 | A = sparse.kron(sparse.eye(B_in), alpha) 569 | # Row-stochastic constraint 570 | R = sparse.kron(sparse.eye(B_in), np.ones(B_out)) 571 | # Formulate constraints 572 | constraints = [ 573 | D @ p <= np.zeros(D.shape[0]), 574 | A @ p == u, 575 | R @ p == np.ones(B_in) 576 | ] 577 | # Build and solve the problem 578 | prob = cvxpy.Problem(objective, constraints) 579 | prob.solve() 580 | if verbose: 581 | print("Solving LP for P given alpha") 582 | print(f"Objective value is {prob.value}") 583 | if prob.value < math.inf: 584 | print(f"Max DP constraint violation is {np.max(constraints[0].violation())}") 585 | print(f"Max unbiased constraint violation is {np.max(constraints[1].violation())}") 586 | print(f"Max row-stochastic constraint violation is {np.max(constraints[2].violation())}") 587 | if prob.value == math.inf: 588 | p.value = np.zeros((B_in * B_out)) 589 | return p.value.reshape((B_in, B_out)), prob.value 590 | 591 | def _solve_qp_for_alpha(self, P, verbose=False): 592 | B_in = 2 ** self.input_bits 593 | B_out = 2 ** self.budget 594 | assert P.shape == (B_in, B_out) 595 | alpha = cvxpy.Variable(B_out) 596 | Q = np.diag(P.sum(axis=0)) 597 | u = np.arange(0, 1+1/B_in, 1/(B_in-1)) 598 | q = -2 * (u.T @ P) 599 | # Impose constraints on ordering (just to make things more easily interpretable) 600 | A = np.zeros((B_out-1, B_out)) 601 | for i in range(B_out-1): 602 | A[i, i] = 1.0 603 | A[i, i+1] = -1.0 604 | objective = cvxpy.Minimize(cvxpy.quad_form(alpha, Q) + q.T @ alpha) 605 | constraints = [P @ alpha == u, A @ alpha <= np.zeros(B_out-1)] 606 | prob = cvxpy.Problem(objective, constraints) 607 | prob.solve() 608 | if verbose: 609 | print("Solving QP for alpha given P") 610 | print(f"Objective value is {prob.value}") 611 | if prob.value < math.inf: 612 | print(f"Max unbiased constraint violation is {np.max(constraints[0].violation())}") 613 | print(f"Max ordering constraint violation is {np.max(constraints[1].violation())}") 614 | return alpha.value, prob.value 615 | 616 | def _run_one_init(self, num_iters, verbose, alphainit, dp_constraint, sparsity): 617 | B_in = 2 ** self.input_bits 618 | B_out = 2 ** self.budget 619 | target = np.arange(0, 1+1/B_in, 1/(B_in-1)) 620 | # Initialize a feasible alpha 621 | alpha = np.linspace(alphainit[0], alphainit[1], num=B_out, endpoint=True) 622 | for iter in range(num_iters): 623 | P, value = self._solve_lp_for_P(alpha, dp_constraint, sparsity, verbose=verbose) 624 | if value < math.inf: 625 | alpha, value = self._solve_qp_for_alpha(P, verbose=verbose) 626 | mse_loss, bias_sq = self.mse_and_bias_squared(P, alpha) 627 | else: 628 | mse_loss = math.inf 629 | bias_sq = math.inf 630 | if verbose: 631 | print("Iteration %d: MSE loss = %.8f, squared bias = %.8f" % ( 632 | iter+1, mse_loss, bias_sq)) 633 | return P, alpha, value 634 | 635 | def _optimize_alt(self, objective="squared", dp_constraint="strict", sparsity=0, num_iters=1, verbose=False, alphainit=None, num_inits=10, Delta=1.0): 636 | if objective != "squared": 637 | raise RuntimeError("Unsupported objective: " + str(objective)) 638 | 639 | best_P = None 640 | best_alpha = None 641 | best_mse_loss = math.inf 642 | if alphainit is None: 643 | # Binary search on initialization, starting from (-Delta, 1+Delta) 644 | delta = Delta 645 | delta_vals = [0.0] 646 | for num_tries in tqdm(range(num_inits), disable=(not verbose)): 647 | delta_vals.append(delta) 648 | alphainit = (0.0 - delta, 1.0 + delta) 649 | if verbose: 650 | print(f"Trying with alphainit={alphainit}") if verbose else None 651 | P, alpha, value = self._run_one_init(num_iters, verbose, alphainit, dp_constraint, sparsity) 652 | if value < math.inf: 653 | mse_loss, bias_sq = self.mse_and_bias_squared(P, alpha) 654 | else: 655 | mse_loss = math.inf 656 | if mse_loss < best_mse_loss: 657 | best_mse_loss = mse_loss 658 | best_P = P 659 | best_alpha = alpha 660 | np_delta_vals = np.array(delta_vals) 661 | lower_delta = np.max(np_delta_vals[np_delta_vals < delta]) 662 | delta = (delta + lower_delta) / 2 663 | else: 664 | if np.any(np.array(delta_vals) > delta): 665 | np_delta_vals = np.array(delta_vals) 666 | higher_delta = np.min(np_delta_vals[np_delta_vals > delta]) 667 | delta = (delta + higher_delta) / 2 668 | else: 669 | delta = 2.0 * delta 670 | P, alpha = best_P, best_alpha 671 | else: 672 | # Just run for the one value provided in alphainit 673 | P, alpha, value = self._run_one_init(num_iters, verbose, alphainit, dp_constraint, sparsity) 674 | if value == math.inf: 675 | print(f"Did not find a feasible solution for alphainit={alphainit}") 676 | mse_loss, bias_sq = self.mse_and_bias_squared(P, alpha) 677 | if verbose: 678 | print("Final: MSE loss = %.8f, squared bias = %.8f" % ( 679 | mse_loss, bias_sq)) 680 | return P, alpha 681 | 682 | # ============================================ 683 | # Trust-region method and supporting functions 684 | # ============================================ 685 | def _get_objective(self): 686 | """ 687 | Makes functions to compute the objective, gradient, and Hessian-vector product 688 | """ 689 | B_in = 2 ** self.input_bits 690 | B_out = 2 ** self.budget 691 | target = np.arange(0, 1+1/B_in, 1/(B_in-1)) 692 | 693 | def objective(x): 694 | """ 695 | Objective function to be minimized. 696 | x is a vectorized version of all optimization variables (entries 697 | of P reshaped as a vector, followed by entries of alpha). 698 | """ 699 | P = x[:B_in*B_out].reshape((B_in, B_out)) 700 | alpha = x[B_in*B_out:] 701 | return (P * np.power(target[:,None] - alpha[None,:], 2)).sum() / B_in 702 | 703 | def jac(x): 704 | """ 705 | Gradient of the objective function at x wrt all parameters. 706 | """ 707 | P = x[:B_in*B_out].reshape((B_in,B_out)) 708 | alpha = x[B_in*B_out:] 709 | g = np.zeros(B_in*B_out + B_out) 710 | g[:B_in*B_out] = np.power(target[:, None] - alpha[None, :], 2).reshape(B_in * B_out) 711 | for j in range(B_out): 712 | g[B_in*B_out + j] = -2 * P[:,j].dot(np.arange(B_in)/(B_in-1) - alpha[j]) 713 | return g 714 | 715 | def hessp(x, p): 716 | """ 717 | Function that returns the product of p with the Hessian of the 718 | objective evaluated at x; i.e., H @ p. 719 | 720 | Note: Not currently used, but may be useful in the future. 721 | """ 722 | P = x[:B_in*B_out].reshape((B_in,B_out)) 723 | alpha = x[B_in*B_out:] 724 | # Hessian block for alpha-P cross terms, a B x (B*B) matrix with B*B non-zero elements 725 | row_ind = np.zeros(B_in * B_out) 726 | col_ind = np.zeros(B_in * B_out) 727 | data = np.zeros(B_in * B_out) 728 | next_nz = 0 729 | for j in range(B_out): 730 | for i in range(B_in): 731 | row_ind[next_nz] = j 732 | col_ind[next_nz] = i*B_out + j 733 | data[next_nz] = -2*(i/(B_in-1) - alpha[j]) 734 | next_nz += 1 735 | hess_alpha_P = sparse.csr_matrix((data, (row_ind, col_ind)), shape=(B_out, B_in*B_out)) 736 | # Hessian block for alpha-alpha derivatives, a B x B matrix with B non-zero elements 737 | hess_alpha_alpha = sparse.diags(2 * P.sum(axis=0)) 738 | Hp = np.zeros(B_in*B_out + B_out) 739 | Hp[:B_in*B_out] = hess_alpha_P.T @ p[B_in*B_out:] 740 | Hp[B_in*B_out:] = hess_alpha_P @ p[:B_in*B_out] + hess_alpha_alpha @ p[B_in*B_out:] 741 | return Hp 742 | 743 | def hess(x): 744 | """ 745 | Function that returns the Hessian of the objective evaluated at x. 746 | """ 747 | P = x[:B_in*B_out].reshape((B_in,B_out)) 748 | alpha = x[B_in*B_out:] 749 | # Hessian block for alpha-P cross terms, a B x (B*B) matrix with B*B non-zero elements 750 | row_ind = np.zeros(B_in * B_out) 751 | col_ind = np.zeros(B_in * B_out) 752 | data = np.zeros(B_in * B_out) 753 | next_nz = 0 754 | for j in range(B_out): 755 | for i in range(B_in): 756 | row_ind[next_nz] = j 757 | col_ind[next_nz] = i*B_out + j 758 | data[next_nz] = -2*(i/(B_in-1) - alpha[j]) 759 | next_nz += 1 760 | hess_alpha_P = sparse.csr_matrix((data, (row_ind, col_ind)), shape=(B_out, B_in*B_out)) 761 | # Hessian block for alpha-alpha derivatives, a B x B matrix with B non-zero elements 762 | hess_alpha_alpha = sparse.diags(2 * P.sum(axis=0)) 763 | return sparse.bmat([[sparse.csr_matrix((B_in*B_out, B_in*B_out)), hess_alpha_P.T], [hess_alpha_P, hess_alpha_alpha]]) 764 | 765 | return objective, jac, hess 766 | 767 | 768 | def _get_x0(self, verbose=False, dp_constraint="strict", init_method="random"): 769 | B_in = 2 ** self.input_bits 770 | B_out = 2 ** self.budget 771 | if init_method == "alt": 772 | # Initialize the trust-region method from the solution returned by the Alternative solver. 773 | P, alpha = self._optimize_alt(self.budget, self.epsilon, dp_constraint=dp_constraint, 774 | verbose=verbose, num_inits=20, Delta=1.0) 775 | if verbose: 776 | print("Initializing from the Alternative solution") 777 | mse_loss, bias_sq = self.mse_and_bias_squared(P, alpha) 778 | print(f"Initial MSE is {mse_loss:.8f} and initial bias squared is {bias_sq:.8f}") 779 | x0 = np.hstack([P.reshape(B_in*B_out), alpha]) 780 | elif init_method == "zeros": 781 | # Initialize from the all-zeros vector 782 | if verbose: 783 | print("Initializing from the all zeros vector") 784 | x0 = np.zeros(B_in*B_out + B_out) 785 | elif init_method == "random": 786 | # Random initialization 787 | if verbose: 788 | print("Initializing from a random vector") 789 | x0 = np.random.randn(B_in*B_out + B_out) 790 | elif init_method == "uniform": 791 | # Uniform initialization 792 | if verbose: 793 | print("Initializing with the uniform strategy") 794 | x0 = np.ones(B_in*B_out + B_out) / B_out 795 | x0[B_in*B_out:] = np.arange(B_out) / (B_out-1) 796 | else: 797 | raise RuntimeError("Unrecognized init_method passed to MVUMechanism with method=`trust-region`") 798 | return x0 799 | 800 | def _get_bounds(self): 801 | B_in = 2 ** self.input_bits 802 | B_out = 2 ** self.budget 803 | ub = np.inf * np.ones(B_in*B_out + B_out) 804 | lb = np.zeros(B_in*B_out + B_out) 805 | lb[B_in*B_out:] = -np.inf 806 | return optimize.Bounds(lb, ub) 807 | 808 | def _get_dp_constraint(self, dp_constraint, sparsity): 809 | B_in = 2 ** self.input_bits 810 | B_out = 2 ** self.budget 811 | D = self._get_dp_constraint_matrix(dp_constraint, sparsity) 812 | num_dp_constraints = D.shape[0] 813 | # This constraint gets applied to the full parameter vector; 814 | # pad with zeros to get the right shape 815 | Dext = sparse.hstack([D, sparse.csr_matrix((num_dp_constraints, B_out))]) 816 | return optimize.LinearConstraint(Dext, -np.inf, 0) 817 | 818 | def _get_row_constraint(self): 819 | B_in = 2 ** self.input_bits 820 | B_out = 2 ** self.budget 821 | R = self._get_row_stochastic_constraint_matrix() 822 | Rext = sparse.hstack([R, sparse.csr_matrix((B_in, B_out))]) 823 | return optimize.LinearConstraint(Rext, 1, 1) 824 | 825 | def _get_unbiased_constraint(self): 826 | B_in = 2 ** self.input_bits 827 | B_out = 2 ** self.budget 828 | target = np.arange(0, 1+1/B_in, 1/(B_in-1)) 829 | 830 | def unbiased_constraint_fn(x): 831 | P = x[:B_in*B_out].reshape((B_in,B_out)) 832 | alpha = x[B_in*B_out:] 833 | return P @ alpha - target 834 | 835 | def unbiased_constraint_jac(x): 836 | P = x[:B_in*B_out].reshape((B_in,B_out)) 837 | alpha = x[B_in*B_out:] 838 | # Return a B by (B*B + B) matrix with 2B * B non-zeros 839 | nnz = 2*B_in*B_out 840 | row_ind = np.zeros(nnz) 841 | col_ind = np.zeros(nnz) 842 | data = np.zeros(nnz) 843 | next_nz = 0 844 | for i in range(B_in): 845 | for j in range(B_out): 846 | # \partial c_i / \partial P_{i,j} 847 | row_ind[next_nz] = i 848 | col_ind[next_nz] = i*B_out + j 849 | data[next_nz] = alpha[j] 850 | next_nz += 1 851 | # \ partial c_i \partial alpha_j 852 | row_ind[next_nz] = i 853 | col_ind[next_nz] = B_in*B_out + j 854 | data[next_nz] = P[i,j] 855 | next_nz += 1 856 | return sparse.csr_matrix((data, (row_ind, col_ind)), shape=(B_in, B_in*B_out + B_out)) 857 | 858 | def unbiased_constraint_hess(x, v): 859 | P = x[:B_in*B_out].reshape((B_in,B_out)) 860 | alpha = x[B_in*B_out:] 861 | # Compute alpha-P Hessian block, a B*B by B matrix with B*B non-zeros 862 | nnz = B_in*B_out 863 | row_ind = np.zeros(nnz) 864 | col_ind = np.zeros(nnz) 865 | data = np.zeros(nnz) 866 | next_nz = 0 867 | for i in range(B_in): 868 | for j in range(B_out): 869 | # Entry corresponding to \partial^2 / (\partial P_{i,j} \partial \alpha_j) 870 | row_ind[next_nz] = i*B_out + j 871 | col_ind[next_nz] = j 872 | data[next_nz] = v[i] 873 | next_nz += 1 874 | hess_block = sparse.csr_matrix((data, (row_ind, col_ind)), shape=(B_in*B_out, B_out)) 875 | return sparse.bmat([[sparse.csr_matrix((B_in*B_out, B_in*B_out)), hess_block], [hess_block.T, sparse.csr_matrix((B_out, B_out))]]) 876 | 877 | unbiased_constraint = optimize.NonlinearConstraint( 878 | unbiased_constraint_fn, np.zeros(B_in), np.zeros(B_in), 879 | unbiased_constraint_jac, 880 | unbiased_constraint_hess, 881 | ) 882 | return unbiased_constraint 883 | 884 | def _get_constraints(self, dp_constraint, sparsity): 885 | dp_constraint = self._get_dp_constraint(dp_constraint, sparsity) 886 | row_constraint = self._get_row_constraint() 887 | unbiased_constraint = self._get_unbiased_constraint() 888 | return [dp_constraint, row_constraint, unbiased_constraint] 889 | 890 | def _optimize_tr(self, objective="squared", dp_constraint="strict", sparsity=0, maxiter=5000, verbose=False, init_method="random"): 891 | if objective != "squared": 892 | raise RuntimeError("Unsupported objective: " + str(objective)) 893 | 894 | if verbose: 895 | verbose_level = 3 896 | else: 897 | verbose_level = 0 898 | 899 | # objective, jac, hessp = self._get_objective() 900 | objective, jac, hess = self._get_objective() 901 | x0 = self._get_x0(verbose=verbose, dp_constraint=dp_constraint, init_method=init_method) 902 | bounds = self._get_bounds() 903 | constraints = self._get_constraints(dp_constraint, sparsity) 904 | self.log = TRLoggerCallback(x0, self.budget, self.epsilon, self.input_bits, init_method, dp_constraint) 905 | 906 | result = optimize.minimize( 907 | objective, x0, method="trust-constr", 908 | jac=jac, hess=hess, bounds=bounds, 909 | constraints=constraints, callback=self.log, 910 | options={"verbose": verbose_level, "maxiter": maxiter, "sparse_jacobian": True}, 911 | ) 912 | if verbose: 913 | if result.success: 914 | print(f"Solver succeeded! {result.message}") 915 | else: 916 | # Note: Even when the solver does not succeed, it doesn't mean 917 | # that the solution is necessarily bad. 918 | print(f"Warning: Solver did not succeed. {result.message}") 919 | 920 | B_in = 2 ** self.input_bits 921 | B_out = 2 ** self.budget 922 | P = result.x[:B_in*B_out].reshape((B_in,B_out)) 923 | alpha = result.x[B_in*B_out:] 924 | # Sort alpha values in ascending order 925 | perm = np.argsort(alpha) 926 | alpha = alpha[perm] 927 | P = P[:,perm] 928 | if verbose: 929 | mse_loss, bias_sq = self.mse_and_bias_squared(P, alpha) 930 | print("Final: MSE loss = %.8f, squared bias = %.8f" % ( 931 | mse_loss, bias_sq)) 932 | return P, alpha 933 | 934 | # ================================================================= 935 | # Master function that dispatches to method-specific implementation 936 | # ================================================================= 937 | def optimize(self, method, **kwargs): 938 | if method == "penalized-lp": 939 | return self._optimize_penalized_lp(**kwargs) 940 | elif method == "alt": 941 | return self._optimize_alt(**kwargs) 942 | elif method == "trust-region": 943 | return self._optimize_tr(**kwargs) 944 | else: 945 | raise RuntimeError(f"Unrecognized method `{method}`. Valid options are penalized-lp, alt, and trust-region.") 946 | 947 | 948 | class InterpolatedMVUMechanism: 949 | 950 | def __init__(self, budget, p, alpha): 951 | self.budget = budget 952 | assert len(p) == 2**budget 953 | self.p = p 954 | self.eta1 = np.log(p) 955 | self.eta2 = self.eta1[::-1] 956 | self.alpha = alpha 957 | 958 | def privatize(self, x): 959 | P = softmax((1 - x[:, None]) * self.eta1[None, :] + x[:, None] * self.eta2[None, :], axis=1) 960 | z = np.array([np.random.choice(P.shape[1], p=P[i]) for i in range(P.shape[0])]) 961 | return z 962 | 963 | def decode(self, z): 964 | assert z.min() >= 0 and z.max() < 2**self.budget 965 | return self.alpha[z.astype(int)] 966 | 967 | 968 | class TRLoggerCallback: 969 | """ 970 | Helper class used in the trust-region solver. 971 | Logs some metrics at each step of the trust-region method, which may be useful 972 | for tuning the solver and/or debugging. 973 | """ 974 | 975 | def __init__(self, x0, budget, epsilon, input_bits, init_method, dp_constraint): 976 | self.now = datetime.now() 977 | self.budget = budget 978 | self.input_bits = input_bits 979 | self.B_in = 2 ** self.input_bits 980 | self.B_out = 2 ** self.budget 981 | self.epsilon = epsilon 982 | self.init_method = init_method 983 | self.dp_constraint = dp_constraint 984 | self.columns = [ 985 | 'optimality', # Infinity norm of the Lagrangian gradient 986 | 'constr_violation', # Maximum constraint violation 987 | 'fun', # Function value 988 | 'tr_radius', # Trust region radius 989 | 'constr_penalty', # Constraint penalty parameter 990 | 'barrier_tolerance', # Tolerance for barrier subproblem 991 | 'barrier_parameter', # Barrier parameter 992 | 'execution_time', # Total execution time 993 | ] 994 | self.results = {} 995 | for col in self.columns: 996 | self.results[col] = [] 997 | self.results['x_diff_norm'] = [] 998 | self.results['P_diff_l1_norm'] = [] 999 | self.results['alpha_diff_l1_norm'] = [] 1000 | self.x_prev = x0.copy() 1001 | self.results['dist_from_init'] = [] 1002 | self.x0 = x0.copy() 1003 | 1004 | def __call__(self, x, result): 1005 | for col in self.columns: 1006 | self.results[col].append(getattr(result, col)) 1007 | self.results['x_diff_norm'].append(np.linalg.norm(self.x_prev - result.x)) 1008 | self.results['P_diff_l1_norm'].append(np.linalg.norm(self.x_prev[:self.B_in*self.B_out] - result.x[:self.B_in*self.B_out], 1)) 1009 | self.results['alpha_diff_l1_norm'].append(np.linalg.norm(self.x_prev[self.B_in*self.B_out:] - result.x[self.B_in*self.B_out:], 1)) 1010 | np.copyto(self.x_prev, result.x) 1011 | self.results['dist_from_init'].append(np.linalg.norm(self.x0 - result.x)) 1012 | return False 1013 | 1014 | 1015 | def save_mechanism(mechanism, path): 1016 | """ 1017 | Save a mechanism to disk. 1018 | """ 1019 | with open(path, 'wb') as f: 1020 | pickle.dump(mechanism, f) 1021 | 1022 | 1023 | def load_mechanism(path, consolidate=False): 1024 | """ 1025 | Load a mechanism from disk. 1026 | """ 1027 | with open(path, 'rb') as f: 1028 | mechanism = pickle.load(f) 1029 | return mechanism 1030 | -------------------------------------------------------------------------------- /mechanisms_pytorch.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | import torch 10 | import torch.nn.functional as F 11 | import numpy as np 12 | import sys 13 | import math 14 | import scipy.optimize as optimize 15 | from utils import optimal_scaling_integer 16 | from hadamard_transform import pad_to_power_of_2, randomized_hadamard_transform, inverse_randomized_hadamard_transform 17 | 18 | 19 | class SkellamMechanismPyTorch: 20 | ''' 21 | Skellam mechanism from https://arxiv.org/pdf/2110.04995.pdf. 22 | ''' 23 | 24 | def __init__(self, budget, d, norm_bound, mu, device, num_clients=1, s=None): 25 | self.budget = budget 26 | self.d = d 27 | self.expanded_d = int(math.pow(2, math.ceil(math.log2(d)))) 28 | self.norm_bound = norm_bound 29 | self.mu = mu 30 | if s is None: 31 | self.s = self.compute_s(num_clients) 32 | # print("s = %.2f" % self.s) 33 | else: 34 | self.s = s 35 | self.scale = optimal_scaling_integer(self.expanded_d, self.s * norm_bound, math.exp(-0.5), tol=1e-3) 36 | if self.scale == 0: 37 | raise RuntimeError("Did not find suitable scale factor; try increasing communication budget") 38 | self.clip_min = -int(math.pow(2, budget - 1)) 39 | self.clip_max = int(math.pow(2, budget - 1)) - 1 40 | self.device = device 41 | self.seed = None 42 | return 43 | 44 | def compute_s(self, num_clients, k=3, rho=1, DIV_EPSILON=1e-22): 45 | """ 46 | Adapted from https://github.com/google-research/federated/blob/master/distributed_dp/accounting_utils.py 47 | """ 48 | def mod_min(gamma): 49 | var = rho / self.d * (num_clients * self.norm_bound)**2 50 | var += (gamma**2 / 4 + self.mu) * num_clients 51 | return k * math.sqrt(var) 52 | 53 | def gamma_opt_fn(gamma): 54 | return (math.pow(2, self.budget) - 2 * mod_min(gamma) / (gamma + DIV_EPSILON))**2 55 | 56 | gamma_result = optimize.minimize_scalar(gamma_opt_fn) 57 | if not gamma_result.success: 58 | raise ValueError('Cannot compute scaling factor.') 59 | return 1. / gamma_result.x 60 | 61 | def renyi_div(self, alphas, l1_norm_bound=None, l2_norm_bound=None): 62 | """ 63 | Computes Renyi divergence of the Skellam mechanism. 64 | """ 65 | if l2_norm_bound is None: 66 | l2_norm_bound = self.norm_bound 67 | if l1_norm_bound is None: 68 | l1_norm_bound = self.norm_bound * min(math.sqrt(self.expanded_d), self.norm_bound) 69 | epsilons = np.zeros(alphas.shape) 70 | B1 = 3 * l1_norm_bound / (2 * self.s ** 3 * self.mu ** 2) 71 | B2 = 3 * l1_norm_bound / (2 * self.s * self.mu) 72 | for i in range(len(alphas)): 73 | alpha = alphas[i] 74 | epsilon = alpha * self.norm_bound ** 2 / (2 * self.mu) 75 | B3 = (2 * alpha - 1) * self.norm_bound ** 2 / (4 * self.s ** 2 * self.mu ** 2) 76 | epsilons[i] = epsilon + min(B1 + B3, B2) 77 | return epsilons 78 | 79 | def dither(self, x): 80 | k = torch.floor(x).to(self.device) 81 | prob = 1 - (x - k) 82 | while True: 83 | output = k + (torch.rand(k.shape).to(self.device) > prob) 84 | if output.norm() <= self.s * self.norm_bound: 85 | break 86 | return output.long() 87 | 88 | def privatize(self, x, same_rotation_batch=False): 89 | assert torch.all(x.norm(2, 1) <= self.norm_bound * (1 + 1e-4)) # add some margin due to clipping rounding issues 90 | assert x.size(1) == self.d 91 | prng = torch.Generator(device=self.device) 92 | self.seed = prng.seed() 93 | x = randomized_hadamard_transform(pad_to_power_of_2(x), prng.manual_seed(self.seed), same_rotation_batch) 94 | z = torch.zeros(x.size()).long().to(self.device) 95 | for i in range(x.shape[0]): 96 | z[i] = self.dither(self.s * x[i]) 97 | scale = self.s**2 * self.mu * torch.ones_like(z) 98 | z += (torch.poisson(scale) - torch.poisson(scale)).long() 99 | z = torch.remainder(z - self.clip_min, self.clip_max - self.clip_min) + self.clip_min 100 | return z 101 | 102 | def decode(self, z, same_rotation_batch=False): 103 | assert self.seed is not None, "Must call privatize before decode." 104 | prng = torch.Generator(device=self.device) 105 | x = inverse_randomized_hadamard_transform(z.float(), prng.manual_seed(self.seed), same_rotation_batch) / self.s 106 | self.seed = None 107 | return x[:, :self.d] 108 | 109 | 110 | class MVUMechanismPyTorch: 111 | 112 | def __init__(self, input_bits, budget, epsilon, P, alpha, norm_bound, device): 113 | self.budget = budget 114 | self.scale = 1 115 | self.epsilon = epsilon 116 | self.input_bits = input_bits 117 | self.P = P.to(device) 118 | self.alpha = alpha.to(device) 119 | self.norm_bound = norm_bound 120 | self.device = device 121 | return 122 | 123 | def dither(self, x): 124 | assert torch.all(x >= 0) and torch.all(x <= 1) 125 | B = 2 ** self.input_bits 126 | k = torch.floor((B-1) * x).to(self.device) 127 | prob = 1 - (B-1) * (x - k/(B-1)) 128 | while True: 129 | output = k + (torch.rand(k.shape).to(self.device) > prob) 130 | if (output / (B-1) - 0.5).norm() <= self.norm_bound: 131 | break 132 | return output.long() 133 | 134 | def privatize(self, x): 135 | z = x.new_zeros(x.size()).long() 136 | for i in range(x.size(0)): 137 | z[i] = self.dither(x[i]) 138 | z = z.flatten() 139 | B = 2 ** self.input_bits 140 | range_B = torch.arange(B).long().to(self.device) 141 | output = torch.zeros(z.shape).long().to(self.device) 142 | for i in range(len(range_B)): 143 | mask = z.eq(range_B[i]) 144 | if mask.sum() > 0: 145 | output[mask] = torch.multinomial(self.P[i], mask.sum(), replacement=True) 146 | return output 147 | 148 | def decode(self, k): 149 | assert k.min() >= 0 and k.max() < 2 ** self.budget 150 | return self.alpha[k] 151 | 152 | 153 | class IMVUMechanismPyTorch: 154 | 155 | def __init__(self, input_bits, budget, P, alpha, device): 156 | self.input_bits = input_bits 157 | self.budget = budget 158 | self.scale = 1 159 | self.P = P.to(device) 160 | self.eta = self.P.log() 161 | self.alpha = alpha.to(device) 162 | self.device = device 163 | return 164 | 165 | def get_etas(self, x): 166 | assert torch.all(x >= 0) and torch.all(x <= 1) 167 | B = 2**self.input_bits 168 | k = torch.floor((B-1) * x).long() 169 | coeff1 = k + 1 - x * (B-1) 170 | coeff2 = x * (B-1) - k 171 | eta1 = self.eta[k] 172 | eta2 = self.eta[k+(coeff2>0)] 173 | return coeff1[:, None] * eta1 + coeff2[:, None] * eta2 174 | 175 | def privatize(self, x, batch_size=int(1e8)): 176 | x = x.flatten() 177 | output = [] 178 | num_batch = int(math.ceil(len(x) / float(batch_size))) 179 | for i in range(num_batch): 180 | lower = i * batch_size 181 | upper = min((i+1) * batch_size, len(x)) 182 | z = x[lower:upper] 183 | P = F.softmax(self.get_etas(z), dim=1) 184 | if P.size(1) == 2: 185 | output.append(torch.bernoulli(P[:, 1])) 186 | else: 187 | output.append(torch.multinomial(P, 1, replacement=True).squeeze()) 188 | output = torch.cat(output, 0).long().to(self.device) 189 | return output 190 | 191 | def decode(self, k): 192 | assert k.min() >= 0 and k.max() < 2 ** self.budget 193 | return self.alpha[k] 194 | 195 | -------------------------------------------------------------------------------- /optimize_mvu.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | import numpy as np 10 | from mechanisms import * 11 | import pickle 12 | import os 13 | import argparse 14 | 15 | if __name__ == "__main__": 16 | parser = argparse.ArgumentParser(description="Optimize and save MVU mechanisms.") 17 | parser.add_argument( 18 | "--method", 19 | default="tr", 20 | type=str, 21 | choices=["tr", "penalized"], 22 | help="which optimization method to use", 23 | ) 24 | parser.add_argument( 25 | "--epsilon", 26 | default=1, 27 | type=float, 28 | help="LDP epsilon", 29 | ) 30 | parser.add_argument( 31 | "--budget", 32 | default=3, 33 | type=int, 34 | help="budget for the MVU mechanism", 35 | ) 36 | parser.add_argument( 37 | "--input_bits", 38 | default=3, 39 | type=int, 40 | help="number of input bits for the MVU mechanism", 41 | ) 42 | parser.add_argument( 43 | "--lam", 44 | default=100, 45 | type=float, 46 | help="soft penalty for penalized solver" 47 | ) 48 | parser.add_argument( 49 | "--num_iters", 50 | default=5, 51 | type=int, 52 | help="number of iterations for penalized solver" 53 | ) 54 | parser.add_argument( 55 | "--dp_constraint", 56 | default="strict", 57 | type=str, 58 | choices=["strict", "metric-l1", "metric-l2"], 59 | help="type of DP constraint" 60 | ) 61 | args = parser.parse_args() 62 | 63 | if args.method == "penalized": 64 | savedir = f"./sweep_eps_budget_penalized_lam{args.lam:0.1e}/" 65 | else: 66 | savedir = f"./sweep_eps_budget_tr/" 67 | os.makedirs(savedir, exist_ok=True) 68 | result_filename = os.path.join( 69 | savedir, f"mechanism_bin{args.input_bits}_bout{args.budget}_{args.dp_constraint}_eps{args.epsilon:0.2f}.pkl") 70 | 71 | if os.path.exists(result_filename): 72 | print(f"{result_filename} already exists, skipping") 73 | exit() 74 | 75 | if args.method == "penalized": 76 | mechanism = MVUMechanism(args.budget, args.epsilon, args.input_bits, method="penalized-lp", verbose=True, 77 | dp_constraint=args.dp_constraint, lam=args.lam, num_iters=args.num_iters) 78 | else: 79 | mechanism = MVUMechanism(args.budget, args.epsilon, args.input_bits, method="trust-region", verbose=True, 80 | dp_constraint=args.dp_constraint, init_method="uniform") 81 | with open(result_filename, "wb") as f: 82 | pickle.dump(mechanism, f) 83 | 84 | -------------------------------------------------------------------------------- /patches/private_prediction.patch: -------------------------------------------------------------------------------- 1 | diff --git a/private_prediction.py b/private_prediction.py 2 | index b17a549..3d55695 100644 3 | --- a/private_prediction.py 4 | +++ b/private_prediction.py 5 | @@ -79,11 +79,11 @@ def get_b_function(epsilon, delta, supremum=True): 6 | gaussian = torch.distributions.normal.Normal(0, 1) 7 | 8 | def b_function(v): 9 | - term = math.exp(epsilon) * gaussian.cdf(-math.sqrt(epsilon * (v + 2))) 10 | + term = math.exp(epsilon) * gaussian.cdf(-math.sqrt(epsilon * (v + 2)) * torch.ones(1)) 11 | if supremum: 12 | - return gaussian.cdf(math.sqrt(epsilon * v)) - term 13 | + return gaussian.cdf(math.sqrt(epsilon * v) * torch.ones(1)) - term 14 | else: 15 | - return -gaussian.cdf(-math.sqrt(epsilon * v)) + term 16 | + return -gaussian.cdf(-math.sqrt(epsilon * v) * torch.ones(1)) + term 17 | 18 | return b_function 19 | 20 | @@ -119,7 +119,7 @@ def sensitivity_scale(epsilon, delta, weight_decay, 21 | 22 | # compute delta knot: 23 | gaussian = torch.distributions.normal.Normal(0, 1) 24 | - delta0 = gaussian.cdf(0) - math.exp(epsilon) * gaussian.cdf(-math.sqrt(2. * epsilon)) 25 | + delta0 = gaussian.cdf(torch.zeros(1)) - math.exp(epsilon) * gaussian.cdf(-math.sqrt(2. * epsilon) * torch.ones(1)) 26 | 27 | # define B-function: 28 | supremum = (delta >= delta0) 29 | @@ -157,7 +157,7 @@ def sensitivity_scale(epsilon, delta, weight_decay, 30 | elif isinstance(criterion, nn.BCELoss): 31 | k = 1.0 32 | else: 33 | - raise ValueError("Lipschitz constant of loss unknown.") 34 | + pass 35 | 36 | # compute final sensitivity scale: 37 | if chaudhuri: 38 | -------------------------------------------------------------------------------- /plot_dme_l1.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | import numpy as np 10 | import matplotlib.pyplot as plt 11 | import seaborn as sns 12 | 13 | epsilons = np.arange(0.5, 5.1, 0.5) 14 | mean_cldp, std_cldp = np.zeros(epsilons.shape), np.zeros(epsilons.shape) 15 | mean_skellam, std_skellam = np.zeros(epsilons.shape), np.zeros(epsilons.shape) 16 | mean_mvu, std_mvu = np.zeros(epsilons.shape), np.zeros(epsilons.shape) 17 | mean_baseline, std_baseline = np.zeros(epsilons.shape), np.zeros(epsilons.shape) 18 | 19 | for i in range(len(epsilons)): 20 | checkpoint = np.load("dme_results/dme_multi_l1_d_128_samples_10000_eps_%.2f_skellam_16_100.00_mvu_3_9.npz" % epsilons[i]) 21 | mean_cldp[i] = checkpoint["squared_error_cldp"].mean() 22 | std_cldp[i] = checkpoint["squared_error_cldp"].std() 23 | mean_skellam[i] = checkpoint["squared_error_skellam"].mean() 24 | std_skellam[i] = checkpoint["squared_error_skellam"].std() 25 | mean_mvu[i] = checkpoint["squared_error_mvu"].mean() 26 | std_mvu[i] = checkpoint["squared_error_mvu"].std() 27 | mean_baseline[i] = checkpoint["squared_error_baseline"].mean() 28 | std_baseline[i] = checkpoint["squared_error_baseline"].std() 29 | 30 | plt.figure(figsize=(8,6)) 31 | with sns.color_palette("deep"): 32 | plt.errorbar(epsilons, mean_cldp, std_cldp, label="CLDP", linewidth=3) 33 | plt.errorbar(epsilons, mean_skellam, std_skellam, label="Skellam ($\\delta > 0$)", linewidth=3) 34 | plt.errorbar(epsilons, mean_mvu, std_mvu, label="MVU", linewidth=3) 35 | plt.errorbar(epsilons, mean_baseline, std_baseline, label="Laplace", color='k', linestyle='--', linewidth=3) 36 | 37 | plt.yscale('log') 38 | plt.xticks(fontsize=20) 39 | plt.xlabel("$\\epsilon$", fontsize=24) 40 | plt.yticks(fontsize=20) 41 | plt.ylabel("MSE", fontsize=24) 42 | plt.legend(fontsize=20, loc="upper right") 43 | plt.grid('on') 44 | plt.title("$L_1, d=128$", fontsize=24) 45 | plt.savefig("figures/dme_l1.pdf", bbox_inches="tight") -------------------------------------------------------------------------------- /plot_dme_l2.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | import numpy as np 10 | import matplotlib.pyplot as plt 11 | import seaborn as sns 12 | from opacus.accountants.analysis.rdp import get_privacy_spent 13 | 14 | epsilons = np.arange(0.5, 5.1, 0.5) 15 | epsilons_approx = np.zeros(epsilons.shape) 16 | mean_cldp, std_cldp = np.zeros(epsilons.shape), np.zeros(epsilons.shape) 17 | mean_skellam, std_skellam = np.zeros(epsilons.shape), np.zeros(epsilons.shape) 18 | mean_mvu, std_mvu = np.zeros(epsilons.shape), np.zeros(epsilons.shape) 19 | mean_mvu_approx, std_mvu_approx = np.zeros(epsilons_approx.shape), np.zeros(epsilons_approx.shape) 20 | mean_baseline, std_baseline = np.zeros(epsilons.shape), np.zeros(epsilons.shape) 21 | 22 | num_samples = 10000 23 | for i in range(len(epsilons)): 24 | checkpoint = np.load("dme_results/dme_multi_l2_d_128_samples_%d_eps_%.2f_skellam_16_15.00_mvu_3_5.npz" % (num_samples, epsilons[i])) 25 | mean_cldp[i] = checkpoint["squared_error_cldp"].mean() 26 | std_cldp[i] = checkpoint["squared_error_cldp"].std() 27 | mean_skellam[i] = checkpoint["squared_error_skellam"].mean() 28 | std_skellam[i] = checkpoint["squared_error_skellam"].std() 29 | mean_mvu[i] = checkpoint["squared_error_mvu"].mean() 30 | std_mvu[i] = checkpoint["squared_error_mvu"].std() 31 | mean_mvu_approx[i] = checkpoint["squared_error_mvu_approx"].mean() 32 | std_mvu_approx[i] = checkpoint["squared_error_mvu_approx"].std() 33 | # use computed metric L2 renyi divergence bounds 34 | renyi_divs = np.load("sweep_eps_budget_penalized_lam1.0e+02/renyi_div_bin5_bout3_metric-l1_eps%.2f_Delta0.50.npz" % (0.5 * epsilons[i])) 35 | epsilons_approx[i] = get_privacy_spent(orders=renyi_divs["alphas"], rdp=renyi_divs["renyi_div_bound"], 36 | delta=(1/(num_samples+1)))[0] 37 | mean_baseline[i] = checkpoint["squared_error_baseline"].mean() 38 | std_baseline[i] = checkpoint["squared_error_baseline"].std() 39 | 40 | plt.figure(figsize=(8,6)) 41 | with sns.color_palette("deep"): 42 | plt.errorbar(epsilons, mean_cldp, std_cldp, label="CLDP", linewidth=3) 43 | plt.errorbar(epsilons, mean_skellam, std_skellam, label="Skellam ($\\delta > 0$)", linewidth=3) 44 | plt.errorbar(epsilons, mean_mvu, std_mvu, label="MVU", linewidth=3) 45 | plt.errorbar(epsilons_approx, mean_mvu_approx, std_mvu_approx, label="MVU ($\\delta > 0$)", linewidth=3, color="lightgreen") 46 | plt.errorbar(epsilons, mean_baseline, std_baseline, label="Gaussian ($\\delta > 0$)", color='k', linestyle='--', linewidth=3) 47 | 48 | plt.yscale('log') 49 | plt.xticks(fontsize=20) 50 | plt.xlabel("$\\epsilon$", fontsize=24) 51 | plt.yticks(fontsize=20) 52 | plt.ylabel("MSE", fontsize=24) 53 | plt.legend(fontsize=20, loc="upper right") 54 | plt.grid('on') 55 | plt.title("$L_2, d=128$", fontsize=24) 56 | plt.savefig("figures/dme_l2.pdf", bbox_inches="tight") 57 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | scipy 2 | numpy 3 | cvxpy 4 | opacus>=1.0 5 | pytorch>=1.7.0 6 | matplotlib 7 | sklearn 8 | pandas 9 | tqdm 10 | seaborn 11 | 12 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | from __future__ import print_function 10 | import argparse 11 | import os 12 | import math 13 | import numpy as np 14 | import torch 15 | import torch.nn as nn 16 | import torch.nn.functional as F 17 | import torch.optim as optim 18 | from torchvision import datasets, transforms, models 19 | from opacus.grad_sample.grad_sample_module import GradSampleModule 20 | from opacus.accountants.analysis.rdp import compute_rdp, get_privacy_spent 21 | from opacus.validators import ModuleValidator 22 | from mechanisms import * 23 | from mechanisms_pytorch import * 24 | from utils import renyi_div_bound_lp, max_divergence_bound, fisher_information_bound, consolidate, params_to_vec, set_grad_to_vec 25 | 26 | import sys 27 | sys.path.append("Handcrafted-DP/") 28 | from data import get_scatter_transform, get_scattered_loader 29 | from models import ScatterLinear 30 | from tqdm import trange 31 | 32 | 33 | class ConvNet(nn.Module): 34 | def __init__(self): 35 | super(ConvNet, self).__init__() 36 | self.conv1 = nn.Conv2d(1, 16, 8, 2, padding=2) 37 | self.conv2 = nn.Conv2d(16, 32, 4, 2, padding=0) 38 | self.fc1 = nn.Linear(32 * 5 * 5, 32) 39 | self.fc2 = nn.Linear(32, 10) 40 | 41 | def forward(self, x): 42 | x = x.view(-1, 1, 28, 28) 43 | x = self.conv1(x) 44 | x = torch.tanh(x) 45 | x = F.avg_pool2d(x, 1) 46 | x = self.conv2(x) 47 | x = torch.tanh(x) 48 | x = F.avg_pool2d(x, 1) 49 | x = torch.flatten(x, 1) 50 | x = self.fc1(x) 51 | x = torch.tanh(x) 52 | x = self.fc2(x) 53 | return x 54 | 55 | 56 | def clip_gradient(args, grad_vec, p=2): 57 | """ 58 | L2 norm clip to args.norm_clip and then L-inf norm clip to args.linf_clip. 59 | """ 60 | C = args.norm_clip 61 | grad_norm = grad_vec.norm(p, 1) 62 | multiplier = grad_norm.new(grad_norm.size()).fill_(1) 63 | multiplier[grad_norm.gt(C)] = C / grad_norm[grad_norm.gt(C)] 64 | grad_vec *= multiplier.unsqueeze(1) 65 | grad_vec.clamp_(-args.linf_clip, args.linf_clip) 66 | return grad_vec 67 | 68 | 69 | def add_noise(args, grad_vec, device, mechanism="gaussian"): 70 | """ 71 | Add noise and quantize the output if args.quantization > 0. 72 | """ 73 | batch_size = grad_vec.size(0) 74 | d = grad_vec.size(1) 75 | if mechanism == "laplace": 76 | dist = torch.distributions.laplace.Laplace(0, 1) 77 | grad_vec += dist.sample(grad_vec.size()).to(device) * args.norm_clip * args.scale 78 | if args.quantization > 0: 79 | assert args.quantization == 1, "Laplace mechanism with quantization level > 1 is not implemented yet." 80 | grad_vec = grad_vec.sign() 81 | elif mechanism == "gaussian": 82 | grad_vec += torch.randn_like(grad_vec).to(device) * args.norm_clip * args.scale 83 | if args.quantization > 0: 84 | assert args.quantization == 1, "Gaussian mechanism with quantization level > 1 is not implemented yet." 85 | grad_vec = grad_vec.sign() 86 | elif isinstance(mechanism, SkellamMechanismPyTorch): 87 | grad_vec = mechanism.decode(mechanism.privatize(mechanism.scale * grad_vec)) 88 | elif isinstance(mechanism, MVUMechanismPyTorch) or isinstance(mechanism, IMVUMechanismPyTorch): 89 | M = args.linf_clip 90 | # scale input to [0,1] 91 | normalized_grad_vec = ((mechanism.scale * grad_vec + M) / (2 * M)).clamp(0, 1) 92 | privatized_grad_vec = mechanism.decode(mechanism.privatize(normalized_grad_vec)) 93 | # scale input back to [-M, M] 94 | grad_vec = privatized_grad_vec.float().view(batch_size, -1) * 2 * M - M 95 | grad_vec /= mechanism.scale 96 | else: 97 | raise NotImplementedError(mechanism) 98 | return grad_vec.sum(0) 99 | 100 | 101 | def train(args, model, device, train_loader, optimizer, epoch, mechanism="gaussian"): 102 | num_param = sum([np.prod(layer.size()) for layer in model.parameters()]) 103 | model.train() 104 | p = 1 if args.mechanism == "laplace" or args.mechanism == "mvu_l1" else 2 105 | for batch_idx, (data, target) in enumerate(train_loader): 106 | data, target = data.to(device), target.to(device) 107 | num_batches = int(math.ceil(float(data.size(0)) / args.physical_batch_size)) 108 | grad_sum = torch.zeros(num_param).to(device) 109 | for i in range(num_batches): 110 | model.zero_grad() 111 | lower = i * args.physical_batch_size 112 | upper = min((i+1) * args.physical_batch_size, data.size(0)) 113 | x, y = data[lower:upper], target[lower:upper] 114 | output = model(x) 115 | loss = F.cross_entropy(output, y) 116 | loss.backward() 117 | grad_vec = params_to_vec(model, return_type="grad_sample") 118 | d = grad_vec.size(1) 119 | if args.norm_clip > 0: 120 | grad_vec = clip_gradient(args, grad_vec, p) 121 | grad_sum += add_noise(args, grad_vec, device, mechanism) 122 | else: 123 | grad_sum += grad_vec.sum(0) 124 | grad_mean = grad_sum / data.size(0) 125 | set_grad_to_vec(model, grad_mean) 126 | optimizer.step() 127 | if batch_idx % args.log_interval == 0: 128 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( 129 | epoch, batch_idx * len(data), len(train_loader.dataset), 130 | 100. * batch_idx / len(train_loader), loss.item())) 131 | 132 | 133 | def test(model, device, test_loader): 134 | model.eval() 135 | test_loss = 0 136 | correct = 0 137 | with torch.no_grad(): 138 | for data, target in test_loader: 139 | data, target = data.to(device), target.to(device) 140 | output = model(data) 141 | test_loss += F.cross_entropy(output, target, reduction='sum').item() # sum up batch loss 142 | pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability 143 | correct += pred.eq(target.view_as(pred)).sum().item() 144 | 145 | test_loss /= len(test_loader.dataset) 146 | 147 | print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( 148 | test_loss, correct, len(test_loader.dataset), 149 | 100. * correct / len(test_loader.dataset))) 150 | return float(correct) / len(test_loader.dataset) 151 | 152 | 153 | def main(): 154 | parser = argparse.ArgumentParser(description='DP-SGD MNIST and CIFAR10 training') 155 | parser.add_argument('--save-dir', type=str, default='dpsgd_results', 156 | help='save directory') 157 | parser.add_argument('--batch-size', type=int, default=600, 158 | help='(virtual) input batch size for training') 159 | parser.add_argument('--physical-batch-size', type=int, default=0, 160 | help='actual batch size') 161 | parser.add_argument('--test-batch-size', type=int, default=1000, 162 | help='input batch size for testing') 163 | parser.add_argument('--dataset', type=str, default='mnist', choices=["mnist", "fmnist", "kmnist", "cifar10"], 164 | help='which dataset to train on') 165 | parser.add_argument('--model', type=str, choices=['linear', 'convnet'], default='convnet', 166 | help='which model to use') 167 | parser.add_argument('--epochs', type=int, default=50, 168 | help='number of epochs to train') 169 | parser.add_argument('--lr', type=float, default=0.1, 170 | help='learning rate') 171 | parser.add_argument('--momentum', type=float, default=0.5, 172 | help='momentum') 173 | parser.add_argument('--norm-clip', type=float, default=0, 174 | help='gradient norm clip') 175 | parser.add_argument('--beta', type=float, default=1, 176 | help='beta scaling for MVU; must be >0') 177 | parser.add_argument('--mechanism', type=str, default='gaussian', 178 | choices=["laplace", "gaussian", "mvu", "mvu_l1", "mvu_l2", "skellam"], 179 | help='which mechanism to use') 180 | parser.add_argument('--quantization', type=int, default=0, 181 | help='quantization level for linf clipping') 182 | parser.add_argument('--input-bits', type=int, default=1, 183 | help='number of input bits for MVU mechanism') 184 | parser.add_argument('--epsilon', type=float, default=1, 185 | help='DP epsilon for MVU mechanism') 186 | parser.add_argument('--scale', type=float, default=0, 187 | help='Laplace/Gaussian noise multiplier') 188 | parser.add_argument('--no-cuda', action='store_true', default=False, 189 | help='disables CUDA training') 190 | parser.add_argument('--log-interval', type=int, default=20, 191 | help='how many batches to wait before logging training status') 192 | parser.add_argument('--save-model', action='store_true', default=False, 193 | help='for saving the current model') 194 | args = parser.parse_args() 195 | os.makedirs(args.save_dir, exist_ok=True) 196 | 197 | use_cuda = not args.no_cuda and torch.cuda.is_available() 198 | args.linf_clip = args.norm_clip / args.beta 199 | if args.physical_batch_size == 0: 200 | args.physical_batch_size = args.batch_size 201 | 202 | device = torch.device("cuda" if use_cuda else "cpu") 203 | 204 | kwargs = {'batch_size': args.batch_size} 205 | if use_cuda: 206 | kwargs.update({'num_workers': 1, 207 | 'pin_memory': True, 208 | 'shuffle': True}, 209 | ) 210 | 211 | if args.mechanism.startswith("mvu"): 212 | output_file = "%s/%s_%s_epochs_%d_lr_%.2e_clip_%.2e_beta_%.2e_%s_bin_%d_quant_%d_eps_%.2e.pth" % ( 213 | args.save_dir, args.dataset, args.model, args.epochs, args.lr, args.norm_clip, args.beta, 214 | args.mechanism, args.input_bits, args.quantization, args.epsilon 215 | ) 216 | else: 217 | output_file = "%s/%s_%s_epochs_%d_lr_%.2e_clip_%.2e_beta_%.2e_%s_quant_%d_scale_%.2e.pth" % ( 218 | args.save_dir, args.dataset, args.model, args.epochs, args.lr, args.norm_clip, args.beta, 219 | args.mechanism, args.quantization, args.scale 220 | ) 221 | 222 | ### DATA LOADING ### 223 | transform = transforms.ToTensor() 224 | if args.dataset == 'mnist': 225 | train_set = datasets.MNIST('./data', train=True, download=True, transform=transform) 226 | test_set = datasets.MNIST('./data', train=False, transform=transform) 227 | elif args.dataset == 'fmnist': 228 | train_set = datasets.FashionMNIST('./data', train=True, download=True, transform=transform) 229 | test_set = datasets.FashionMNIST('./data', train=False, transform=transform) 230 | elif args.dataset == 'kmnist': 231 | train_set = datasets.KMNIST('./data', train=True, download=True, transform=transform) 232 | test_set = datasets.KMNIST('./data', train=False, transform=transform) 233 | else: 234 | transform = transforms.Compose([ 235 | transforms.Resize((32, 32)), 236 | transforms.ToTensor(), 237 | transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), 238 | ]) 239 | train_set = datasets.CIFAR10( 240 | root='../data', train=True, download=True, transform=transform) 241 | test_set = datasets.CIFAR10( 242 | root='../data', train=False, download=True, transform=transform) 243 | 244 | train_loader = torch.utils.data.DataLoader( 245 | train_set, batch_size=args.batch_size, shuffle=True, num_workers=1, pin_memory=True) 246 | test_loader = torch.utils.data.DataLoader( 247 | test_set, batch_size=args.batch_size, shuffle=False, num_workers=1, pin_memory=True) 248 | 249 | ### MODEL LOADING ### 250 | if args.model == "linear": 251 | scattering, K, (h, w) = get_scatter_transform("cifar10" if args.dataset == "cifar10" else "mnist") 252 | scattering.to(device) 253 | train_loader = get_scattered_loader(train_loader, scattering, device, drop_last=True, sample_batches=False) 254 | test_loader = get_scattered_loader(test_loader, scattering, device) 255 | model = GradSampleModule(ScatterLinear(K, (h, w), input_norm="GroupNorm", num_groups=27).to(device)) 256 | else: 257 | assert args.dataset != "cifar10", "CIFAR-10 ConvNet training is not supported." 258 | model = GradSampleModule(ConvNet().to(device)) 259 | num_param = sum([np.prod(layer.size()) for layer in model.parameters()]) 260 | print("Number of model parameters = %d" % num_param) 261 | optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) 262 | 263 | ### MECHANISM LOADING ### 264 | if args.mechanism.startswith("mvu") and args.norm_clip > 0 and args.linf_clip > 0: 265 | savefile = os.path.join('sweep_eps_budget_penalized_lam1.0e+02', f"mechanism_bin{args.input_bits}_bout{args.quantization}_metric-l1_eps{args.epsilon:.2f}.pkl") 266 | with open(savefile, "rb") as file: 267 | mechanism_numpy = pickle.load(file) 268 | mechanism_numpy.P /= mechanism_numpy.P.sum(1)[:, None] 269 | norm_bound = args.norm_clip / (2 * args.linf_clip) # L2 sensitivity for leave-one-out adjacency 270 | if args.mechanism == "mvu": 271 | mechanism = MVUMechanismPyTorch( 272 | args.input_bits, args.quantization, args.epsilon, torch.from_numpy(mechanism_numpy.P), 273 | torch.from_numpy(mechanism_numpy.alpha), norm_bound, device) 274 | mechanism.scale = 0.9 275 | print("Computing Renyi divergence bounds using LP relaxation") 276 | renyi_div_bounds = renyi_div_bound_lp(orders, num_param, mechanism_numpy.P, norm_bound, greedy=True) 277 | elif args.mechanism == "mvu_l1": 278 | mechanism = IMVUMechanismPyTorch( 279 | args.input_bits, args.quantization, torch.from_numpy(mechanism_numpy.P), 280 | torch.from_numpy(mechanism_numpy.alpha), device) 281 | log_P = np.log(mechanism_numpy.P) 282 | epsilon_bound = args.epsilon + max_divergence_bound(log_P) 283 | print("Max divergence bound = %.4f" % epsilon_bound) 284 | else: 285 | assert args.input_bits == 1, "Interpolated MVU is not defined for b_in > 1" 286 | mechanism_numpy.P[1, :] = np.flip(mechanism_numpy.P[0, :], (0,)) 287 | mechanism = IMVUMechanismPyTorch( 288 | args.input_bits, args.quantization, torch.from_numpy(mechanism_numpy.P), 289 | torch.from_numpy(mechanism_numpy.alpha), device) 290 | P, _ = consolidate(mechanism_numpy) 291 | fisher_info_bound = fisher_information_bound(P[0, :]) 292 | print("Fisher info bound = %.4f" % fisher_info_bound) 293 | elif args.mechanism == "skellam": 294 | mu = (args.scale * args.norm_clip)**2 295 | mechanism = SkellamMechanismPyTorch(args.quantization, num_param, args.norm_clip, mu, device) 296 | else: 297 | mechanism = args.mechanism 298 | 299 | q = args.batch_size / float(len(train_set)) 300 | orders = np.array(list(np.linspace(1.1, 10.9, 99)) + list(range(11, 64))) 301 | test_accs, epsilons = torch.zeros(args.epochs), torch.zeros(args.epochs) 302 | for epoch in range(1, args.epochs + 1): 303 | ### TRAINING ### 304 | train(args, model, device, train_loader, optimizer, epoch, mechanism) 305 | test_accs[epoch-1] = test(model, device, test_loader) 306 | ### PRIVACY ACCOUNTING ### 307 | delta = 1e-5 308 | if args.mechanism == "laplace" and args.scale > 0 and args.norm_clip > 0: 309 | delta = 0 310 | opt_order = 0 311 | epsilon = epoch / args.scale 312 | elif args.mechanism == "gaussian" and args.scale > 0 and args.norm_clip > 0: 313 | rdp_const = epoch * orders / (2 * args.scale ** 2) 314 | epsilon, opt_order = get_privacy_spent(orders=orders, rdp=rdp_const, delta=delta) 315 | elif args.mechanism == "mvu" and args.norm_clip > 0: 316 | rdp_const = renyi_div_bounds * epoch 317 | epsilon, opt_order = get_privacy_spent(orders=orders, rdp=rdp_const, delta=delta) 318 | elif args.mechanism == "mvu_l1" and args.norm_clip > 0: 319 | delta = 0 320 | opt_order = 0 321 | epsilon = epoch * epsilon_bound * norm_bound 322 | elif args.mechanism == "mvu_l2" and args.norm_clip > 0: 323 | rdp_const = epoch * orders * fisher_info_bound * norm_bound**2 / 2 324 | epsilon, opt_order = get_privacy_spent(orders=orders, rdp=rdp_const, delta=delta) 325 | elif args.mechanism == "skellam" and args.norm_clip > 0: 326 | rdp_const = epoch * mechanism.renyi_div(orders) 327 | epsilon, opt_order = get_privacy_spent(orders=orders, rdp=rdp_const, delta=delta) 328 | else: 329 | epsilon, opt_order = math.inf, 0 330 | epsilons[epoch-1] = epsilon 331 | print("Epsilon at delta=%.2e: %.4f, optimal alpha: %.4f\n" % (delta, epsilon, opt_order)) 332 | 333 | if args.save_model: 334 | if os.path.exists(output_file): 335 | checkpoint = torch.load(output_file) 336 | checkpoint['state_dict'].append(model.state_dict()) 337 | checkpoint['test_accs'].append(test_accs) 338 | checkpoint['epsilons'].append(epsilons) 339 | torch.save(checkpoint, output_file) 340 | else: 341 | torch.save({'state_dict': [model.state_dict()], 'test_accs': [test_accs], 'epsilons': [epsilons]}, output_file) 342 | 343 | 344 | if __name__ == '__main__': 345 | main() 346 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | # 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | import numpy as np 10 | from tqdm import tqdm 11 | from scipy.optimize import linprog 12 | from scipy.special import softmax 13 | import math 14 | import torch 15 | 16 | import sys 17 | sys.path.append("private_prediction/") 18 | from util import binary_search 19 | sys.path.append("fastwht/python/") 20 | from hadamard import * 21 | 22 | 23 | class FWHTRandomProjector: 24 | ''' 25 | Fast random projector using the Fast Walsh-Hadamard Transform. 26 | 27 | Arguments: 28 | num_parameters: Number of input data dimensions 29 | rank: Rank of projection 30 | seed: Random seed to control the projection. 31 | Default None uses a different random seed each time. 32 | ''' 33 | 34 | def __init__(self, num_parameters, expanded_dim): 35 | self.num_parameters = num_parameters 36 | self.expanded_dim = expanded_dim 37 | self.D = np.sign(np.random.normal(0, 1, self.expanded_dim)) 38 | 39 | def project(self, x): 40 | x_expanded = np.zeros(self.expanded_dim) 41 | x_expanded[:len(x)] = x 42 | z = fastwht(x_expanded * self.D, order='hadamard') * math.sqrt(float(x_expanded.shape[0])) 43 | return z 44 | 45 | def inverse(self, x): 46 | out_expanded = fastwht(x, order='hadamard') * self.D * math.sqrt(float(x.shape[0])) 47 | return out_expanded[:self.num_parameters] 48 | 49 | 50 | def _log_add(logx, logy): 51 | a, b = np.minimum(logx, logy), np.maximum(logx, logy) 52 | eq_zero = (a == -np.inf) 53 | # Use exp(a) + exp(b) = (exp(a - b) + 1) * exp(b) 54 | result = np.log(np.exp(logx - logy) + 1) + logy 55 | result[eq_zero] = b[eq_zero] 56 | return result 57 | 58 | 59 | def log_sum(log_zs): 60 | shape = log_zs.sum(-1).shape 61 | log_zs = log_zs.reshape(-1, log_zs.shape[-1]) 62 | result = log_zs[:, 0] 63 | for i in range(1, log_zs.shape[-1]): 64 | result = _log_add(result, log_zs[:, i]) 65 | return result 66 | 67 | 68 | def renyi_div_stable(P, order): 69 | """ 70 | Computes (log) Renyi divergence using stable log-space computation 71 | """ 72 | assert order > 1, "Renyi divergence order must be > 1" 73 | assert np.all(P >= 0) 74 | log_P = np.log(P) 75 | log_Q = np.copy(log_P) 76 | log_Q[log_Q == -np.inf] = 0 77 | log_zs = order * log_P[:, None, :] - (order - 1) * log_Q[None, :, :] 78 | return np.reshape(log_sum(log_zs) / (order - 1), (P.shape[0], P.shape[0])) 79 | 80 | 81 | def renyi_div_bound_lp(alphas, d, P, Delta, greedy=False, verbose=False): 82 | """ 83 | Computes LP relaxation upper bound for Renyi divergence of MVU mechanism. 84 | If greedy=True, always compute greedy solution. 85 | """ 86 | B = P.shape[0] 87 | bounds = np.zeros(alphas.shape) 88 | for i in tqdm(range(alphas.shape[0])): 89 | alpha = alphas[i] 90 | D = renyi_div_stable(P, alpha) 91 | D[D < 0] = 0 92 | irange = np.arange(0, B) / (B-1) 93 | # add identity to cost to avoid division by zero 94 | cost = (np.power(irange[:, None] - irange[None, :], 2) + np.eye(B)).flatten() 95 | opt = (D / cost).argmax() 96 | max_div, cost_per_dim = D[opt], cost[opt] 97 | if d * cost_per_dim < Delta**2 and (not greedy): 98 | c = -np.tile(D, d) 99 | A = np.tile(np.power(irange[:, None] - irange[None, :], 2).flatten(), d)[None, :] 100 | A = np.concatenate([A, np.kron(np.eye(d), np.ones((1, B*B)))], 0) 101 | b = np.array([0.25] + [1] * d) 102 | res = linprog(c, A_ub=A, b_ub=b, options={"disp": verbose}) 103 | bounds[i] = -res.fun 104 | else: 105 | # compute bound using optimal greedy solution 106 | num_items = Delta**2 / cost_per_dim 107 | bounds[i] = max_div * num_items 108 | return bounds 109 | 110 | 111 | def optimal_scaling_mvu(samples, mechanism, conf, p=2): 112 | """ 113 | Computes optimal scaling factor for MVU mechanism to ensure that vector norm after 114 | dithering does not increase. 115 | """ 116 | B = 2 ** mechanism.input_bits 117 | def constraint(t): 118 | x = t * samples 119 | x = np.clip((x+1)/2, 0, 1) 120 | z = mechanism.dither(x, mechanism.input_bits) / (B-1) 121 | z = 2*z - 1 122 | norms = np.linalg.norm(z, p, 1) 123 | return (norms > 1).astype(float).mean() < conf 124 | return binary_search(lambda t: t, constraint, 0, 1, tol=1e-3) 125 | 126 | 127 | def post_rounding_l2_norm_bound(d, l2_norm_bound, beta): 128 | """ 129 | Function for computing vector norm bound after quantizing to the integer grid. 130 | Adapted from https://github.com/google-research/federated/blob/master/distributed_dp/compression_utils.py 131 | """ 132 | bound1 = l2_norm_bound + math.sqrt(d) 133 | squared_bound2 = l2_norm_bound**2 + 0.25 * d 134 | squared_bound2 += (math.sqrt(2.0 * math.log(1.0 / beta)) * (l2_norm_bound + 0.5 * math.sqrt(d))) 135 | bound2 = math.sqrt(squared_bound2) 136 | # bound2 is inf if beta = 0, in which case we fall back to bound1. 137 | return min(bound1, bound2) 138 | 139 | 140 | def optimal_scaling_integer(d, l2_norm_bound, beta, tol=1e-3): 141 | """ 142 | Computes optimal scaling factor for DDG/Skellam mechanism to ensure that vector norm after 143 | dithering does not increase. 144 | """ 145 | def constraint(t): 146 | if t == 0: 147 | return True 148 | quantized_norm = post_rounding_l2_norm_bound(d, t, beta) 149 | return quantized_norm <= l2_norm_bound + 1e-6 150 | opt_norm = binary_search(lambda t: t, constraint, 0, l2_norm_bound, tol=tol) 151 | return opt_norm / l2_norm_bound 152 | 153 | 154 | def max_divergence_bound(log_P, precision=1e-4): 155 | """ 156 | Computes L1 metric DP additive factor for the interpolated MVU mechanism. 157 | """ 158 | B = log_P.shape[0] 159 | epsilon_bound = np.zeros(B - 1) 160 | for i in range(log_P.shape[0] - 1): 161 | theta = log_P[i+1] - log_P[i] 162 | xs = np.arange(0, 1 + precision/2, precision) 163 | etas = xs[:, None] * log_P[None, i+1] + (1 - xs)[:, None] * log_P[None, i] 164 | func = abs(softmax(etas, axis=1) @ theta) * (B - 1) 165 | epsilon_bound[i] = func.max() 166 | return epsilon_bound.max() 167 | 168 | 169 | def fisher_information(bs, eta1, eta2): 170 | """ 171 | Computes Fisher information of the interpolated MVU mechanism for a range of coefficients. 172 | """ 173 | etas = (1 - bs)[:, None] * eta1[None, :] + bs[:, None] * eta2[None, :] 174 | P = softmax(etas, axis=1) 175 | fi = P @ np.power(eta2 - eta1, 2) - np.power(P @ (eta2 - eta1), 2) 176 | return fi 177 | 178 | 179 | def fisher_information_bound(p, precision=1e-4): 180 | """ 181 | Asserts the condition that maximum FI occurs in range [0,1] and returns this maximum. 182 | If condition fails, return -1. 183 | """ 184 | 185 | # sampling probability vectors must be symmetric about 0.5(!) 186 | eta1 = np.log(p) 187 | eta2 = eta1[::-1] 188 | 189 | # compute maximum FI within range [0,1] 190 | bs = np.arange(0, 1+precision/2, precision) 191 | fi_max = fisher_information(bs, eta1, eta2).max() 192 | 193 | # compute range [a_min, a_max] 194 | j_max = np.argmax(eta2 - eta1) 195 | eta_sq_max = np.power(eta2 - eta1, 2)[j_max] 196 | c = fi_max / (4 * eta_sq_max) 197 | p = (1 + math.sqrt(1 - 4*c)) / 2 198 | assert p >= 0.5 199 | a_max = 1 200 | while True: 201 | eta = (1-a_max) * eta1 + a_max * eta2 202 | p_eta = softmax(eta) 203 | if p_eta[j_max] >= p: 204 | break 205 | else: 206 | a_max += 1 207 | a_min = 1 - a_max 208 | 209 | # compute maximum FI within range [a_min, a_max] 210 | bs = np.arange(a_min, a_max + precision/2, precision) 211 | fi = fisher_information(bs, eta1, eta2) 212 | return fi.max() 213 | 214 | 215 | def consolidate(mechanism, tol=1e-8): 216 | """ 217 | Consolidates the sampling probability matrix P to remove redundant columns. 218 | Must be called prior to computing the Fisher information to avoid infinite loop! 219 | """ 220 | digits = int(-math.log10(tol)) 221 | uniques = np.unique(mechanism.alpha.round(digits)) 222 | P_new = [] 223 | alpha_new = [] 224 | for i in range(len(uniques)): 225 | indices = np.arange(0, len(mechanism.alpha))[np.isclose(mechanism.alpha, uniques[i])] 226 | p = mechanism.P[:, indices].sum(1) 227 | P_new.append(p[:, None]) 228 | alpha_new.append(mechanism.alpha[indices[0]]) 229 | P_new = np.concatenate(P_new, 1) 230 | alpha_new = np.hstack(alpha_new) 231 | return P_new, alpha_new 232 | 233 | 234 | def params_to_vec(model, return_type="param"): 235 | ''' 236 | Helper function that concatenates model parameters or gradients into a single vector. 237 | ''' 238 | vec = [] 239 | for param in model.parameters(): 240 | if return_type == "param": 241 | vec.append(param.data.view(1, -1)) 242 | elif return_type == "grad": 243 | vec.append(param.grad.view(1, -1)) 244 | elif return_type == "grad_sample": 245 | if hasattr(param, "grad_sample"): 246 | vec.append(param.grad_sample.view(param.grad_sample.size(0), -1)) 247 | else: 248 | print("Error: Per-sample gradient not found") 249 | sys.exit(1) 250 | return torch.cat(vec, dim=1).squeeze() 251 | 252 | 253 | def set_grad_to_vec(model, vec): 254 | ''' 255 | Helper function that sets the model's gradient to a given vector. 256 | ''' 257 | model.zero_grad() 258 | for param in model.parameters(): 259 | size = param.data.view(1, -1).size(1) 260 | param.grad = vec[:size].view_as(param.data).clone() 261 | vec = vec[size:] 262 | return 263 | 264 | --------------------------------------------------------------------------------